move cosine loss compute to function, fix some args

Author: francoishernandez

onmt/modules/copy_generator.py

@@ -186,22 +186,23 @@ def __init__(self, criterion, generator, tgt_vocab, normalize_by_length,
         self.tgt_vocab = tgt_vocab
         self.normalize_by_length = normalize_by_length
 
-    def _make_shard_state(self, batch, output, range_, attns):
+    def _make_shard_state(self, batch, output, enc_src, enc_tgt, range_, attns):
         """See base class for args description."""
         if getattr(batch, "alignment", None) is None:
             raise AssertionError("using -copy_attn you need to pass in "
                                  "-dynamic_dict during preprocess stage.")
 
         shard_state = super(CopyGeneratorLossCompute, self)._make_shard_state(
-            batch, output, range_, attns)
+            batch, output, enc_src, enc_tgt, range_, attns)
 
         shard_state.update({
             "copy_attn": attns.get("copy"),
             "align": batch.alignment[range_[0] + 1: range_[1]]
         })
         return shard_state
 
-    def _compute_loss(self, batch, output, target, copy_attn, align,
+    def _compute_loss(self, batch, normalization, output, target,
+                      copy_attn, align, enc_src=None, enc_tgt=None,
                       std_attn=None, coverage_attn=None):
         """Compute the loss.
 
@@ -244,8 +245,18 @@ def _compute_loss(self, batch, output, target, copy_attn, align,
         offset_align = align[correct_mask] + len(self.tgt_vocab)
         target_data[correct_mask] += offset_align
 
+        if self.lambda_cosine != 0.0:
+            cosine_loss, num_ex = self._compute_cosine_loss(enc_src, enc_tgt)
+            loss += self.lambda_cosine * (cosine_loss / num_ex)
+        else:
+            cosine_loss = None
+            num_ex = 0
+
         # Compute sum of perplexities for stats
-        stats = self._stats(loss.sum().clone(), scores_data, target_data)
+        stats = self._stats(loss.sum().clone(),
+                            cosine_loss.clone() if cosine_loss is not None
+                            else cosine_loss,
+                            scores_data, target_data, num_ex)
 
         # this part looks like it belongs in CopyGeneratorLoss
         if self.normalize_by_length:

onmt/utils/loss.py

@@ -92,7 +92,8 @@ def __init__(self, criterion, generator):
     def padding_idx(self):
         return self.criterion.ignore_index
 
-    def _make_shard_state(self, batch, output, range_, attns=None):
+    def _make_shard_state(self, batch, enc_src, enc_tgt,
+                          output, range_, attns=None):
         """
         Make shard state dictionary for shards() to return iterable
         shards for efficient loss computation. Subclass must define
@@ -315,14 +316,7 @@ def _compute_loss(self, batch, normalization, output, target,
         loss = loss/float(normalization)
 
         if self.lambda_cosine != 0.0:
-            max_src = enc_src.max(axis=0)[0]
-            max_tgt = enc_tgt.max(axis=0)[0]
-            cosine_loss = torch.nn.functional.cosine_similarity(
-                max_src.float(), max_tgt.float(), dim=1)
-            ones = torch.ones(cosine_loss.size()).to(cosine_loss.device)
-            cosine_loss = ones - cosine_loss
-            num_ex = cosine_loss.size(0)
-            cosine_loss = cosine_loss.sum()
+            cosine_loss, num_ex = self._compute_cosine_loss(enc_src, enc_tgt)
             loss += self.lambda_cosine * (cosine_loss / num_ex)
         else:
             cosine_loss = None
@@ -340,6 +334,16 @@ def _compute_coverage_loss(self, std_attn, coverage_attn):
         covloss *= self.lambda_coverage
         return covloss
 
+    def _compute_cosine_loss(self, enc_src, enc_tgt):
+        max_src = enc_src.max(axis=0)[0]
+        max_tgt = enc_tgt.max(axis=0)[0]
+        cosine_loss = torch.nn.functional.cosine_similarity(
+            max_src.float(), max_tgt.float(), dim=1)
+        ones = torch.ones(cosine_loss.size()).to(cosine_loss.device)
+        cosine_loss = ones - cosine_loss
+        num_ex = cosine_loss.size(0)
+        return cosine_loss.sum(), num_ex
+
     def _compute_alignement_loss(self, align_head, ref_align):
         """Compute loss between 2 partial alignment matrix."""
         # align_head contains value in [0, 1) presenting attn prob,