def metric_fn(y_true, y_pred): """Returns the in_top_k metric.""" softmax_logits = y_pred[0, :] logits = tf.slice(softmax_logits, [0, 1], [batch_size, 1]) # The dup mask should be obtained from input data, but we did not yet find # a good way of getting it with keras, so we set it to zeros to neglect the # repetition correction dup_mask = tf.zeros([batch_size, 1]) _, _, in_top_k, _, _ = ( neumf_model.compute_eval_loss_and_metrics_helper( logits, softmax_logits, dup_mask, params["num_neg"], params["match_mlperf"], params["use_xla_for_gpu"])) is_training = tf.keras.backend.learning_phase() if isinstance(is_training, int): is_training = tf.constant(bool(is_training), dtype=tf.bool) in_top_k = tf.cond( is_training, lambda: tf.zeros(shape=in_top_k.shape, dtype=in_top_k.dtype), lambda: in_top_k) return in_top_k
def metric_fn(y_true, y_pred): """Returns the in_top_k metric.""" softmax_logits = y_pred[0, :] logits = tf.slice(softmax_logits, [0, 1], [batch_size, 1]) # The dup mask should be obtained from input data, but we did not yet find # a good way of getting it with keras, so we set it to zeros to neglect the # repetition correction dup_mask = tf.zeros([batch_size, 1]) _, _, in_top_k, _, _ = ( neumf_model.compute_eval_loss_and_metrics_helper( logits, softmax_logits, dup_mask, params["num_neg"], params["match_mlperf"], params["use_xla_for_gpu"])) is_training = tf.keras.backend.learning_phase() if isinstance(is_training, int): is_training = tf.constant(bool(is_training), dtype=tf.bool) in_top_k = tf.cond( is_training, lambda: tf.zeros(shape=in_top_k.shape, dtype=in_top_k.dtype), lambda: in_top_k) return in_top_k
def _model_fn(features, labels, mode, params): """Model Function for NeuMF estimator.""" logits = logits_fn(features, feature_columns, params) # Softmax with the first column of zeros is equivalent to sigmoid. softmax_logits = tf.concat( [tf.zeros(logits.shape, dtype=logits.dtype), logits], axis=1) if mode == tf.estimator.ModeKeys.EVAL: duplicate_mask = tf.cast(features[rconst.DUPLICATE_MASK], tf.float32) cross_entropy, metric_fn, in_top_k, ndcg, metric_weights = ( neumf_model.compute_eval_loss_and_metrics_helper( logits, softmax_logits, duplicate_mask, params["num_neg"], params["match_mlperf"], use_tpu_spec=params["use_tpu"])) return tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=cross_entropy, eval_metrics=(metric_fn, [in_top_k, ndcg, metric_weights])) elif mode == tf.estimator.ModeKeys.TRAIN: labels = tf.cast(labels, tf.int32) valid_pt_mask = features[rconst.VALID_POINT_MASK] optimizer = tf.train.AdamOptimizer( learning_rate=params["learning_rate"], beta1=params["beta1"], beta2=params["beta2"], epsilon=params["epsilon"]) optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer) loss = tf.losses.sparse_softmax_cross_entropy( labels=labels, logits=softmax_logits, weights=tf.cast(valid_pt_mask, tf.float32)) # This tensor is used by logging hooks. tf.identity(loss, name="cross_entropy") global_step = tf.train.get_global_step() tvars = tf.trainable_variables() gradients = optimizer.compute_gradients( loss, tvars, colocate_gradients_with_ops=True) minimize_op = optimizer.apply_gradients(gradients, global_step=global_step, name="train") update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op = tf.group(minimize_op, update_ops) return tf.contrib.tpu.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op) else: raise NotImplementedError