scale loss by num replicas

5b81bb59 · guptapriya · guptapriya · 346b570f · 5b81bb59
Commit 5b81bb59 authored Jun 05, 2019 by guptapriya Committed by guptapriya Jun 04, 2019
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

official/recommendation/ncf_keras_main.py official/recommendation/ncf_keras_main.py +2 -1

No files found.
--- a/official/recommendation/ncf_keras_main.py
+++ b/official/recommendation/ncf_keras_main.py
@@ -235,10 +235,11 @@ def _get_keras_model(params):
      from_logits=True,
      reduction="sum")

+  loss_scale_factor = batch_size * tf.distribute.get_strategy().num_replicas_in_sync
  keras_model.add_loss(loss_obj(
      y_true=label_input,
      y_pred=softmax_logits,
-      sample_weight=valid_pt_mask_input) * 1.0 / batch_size)
+      sample_weight=valid_pt_mask_input) * 1.0 / loss_scale_factor)

  keras_model.summary()
  return keras_model