disable sharded loss if lambda_cosine

Author: francoishernandez

onmt/utils/loss.py

@@ -167,7 +167,7 @@ def __call__(self,
             return loss, stats
         batch_stats = onmt.utils.Statistics()
         for shard in shards(shard_state, shard_size):
-            loss, stats = self._compute_loss(batch, **shard)
+            loss, stats = self._compute_loss(batch, normalization, **shard)
             loss.backward()
             batch_stats.update(stats)
         return None, batch_stats
@@ -243,9 +243,7 @@ def _make_shard_state(self, batch, output, enc_src, enc_tgt,
                           range_, attns=None):
         shard_state = {
             "output": output,
-            "target": batch.tgt[range_[0] + 1: range_[1], :, 0],
-            "enc_src": enc_src,
-            "enc_tgt": enc_tgt
+            "target": batch.tgt[range_[0] + 1: range_[1], :, 0]
         }
         if self.lambda_coverage != 0.0:
             coverage = attns.get("coverage", None)
@@ -283,10 +281,15 @@ def _make_shard_state(self, batch, output, enc_src, enc_tgt,
                 "align_head": attn_align,
                 "ref_align": ref_align[:, range_[0] + 1: range_[1], :]
             })
+        if self.lambda_cosine != 0.0:
+            shard_state.update({
+                "enc_src": enc_src,
+                "enc_tgt": enc_tgt
+                })
         return shard_state
 
     def _compute_loss(self, batch, normalization, output, target,
-                      enc_src, enc_tgt, std_attn=None,
+                      enc_src=None, enc_tgt=None, std_attn=None,
                       coverage_attn=None, align_head=None, ref_align=None):
 
         bottled_output = self._bottle(output)
@@ -400,7 +403,7 @@ def shards(state, shard_size, eval_only=False):
         # over the shards, not over the keys: therefore, the values need
         # to be re-zipped by shard and then each shard can be paired
         # with the keys.
-        for shard_tensors in zip(*values):
+        for i, shard_tensors in enumerate(zip(*values)):
             yield dict(zip(keys, shard_tensors))
 
         # Assumed backprop'd

onmt/utils/parse.py

@@ -120,6 +120,10 @@ def validate_train_opts(cls, opt):
         assert len(opt.attention_dropout) == len(opt.dropout_steps), \
             "Number of attention_dropout values must match accum_steps values"
 
+        assert not(opt.max_generator_batches > 0 and opt.lambda_cosine != 0), \
+            "-lambda_cosine loss is not implemented for max_generator_batches > 0."
+
+
     @classmethod
     def validate_translate_opts(cls, opt):
         if opt.beam_size != 1 and opt.random_sampling_topk != 1: