Implementing coverage loss of abisee (2017) (#1464)

* Implementing coverage loss of abisee (2017)
* fix lambda_coverage value

Author: pltrdy

onmt/modules/copy_generator.py

@@ -2,7 +2,7 @@
 import torch.nn as nn
 
 from onmt.utils.misc import aeq
-from onmt.utils.loss import LossComputeBase
+from onmt.utils.loss import NMTLossCompute
 
 
 def collapse_copy_scores(scores, batch, tgt_vocab, src_vocabs=None,
@@ -177,10 +177,12 @@ def forward(self, scores, align, target):
         return loss
 
 
-class CopyGeneratorLossCompute(LossComputeBase):
+class CopyGeneratorLossCompute(NMTLossCompute):
     """Copy Generator Loss Computation."""
-    def __init__(self, criterion, generator, tgt_vocab, normalize_by_length):
-        super(CopyGeneratorLossCompute, self).__init__(criterion, generator)
+    def __init__(self, criterion, generator, tgt_vocab, normalize_by_length,
+                 lambda_coverage=0.0):
+        super(CopyGeneratorLossCompute, self).__init__(
+            criterion, generator, lambda_coverage=lambda_coverage)
         self.tgt_vocab = tgt_vocab
         self.normalize_by_length = normalize_by_length
 
@@ -190,14 +192,17 @@ def _make_shard_state(self, batch, output, range_, attns):
             raise AssertionError("using -copy_attn you need to pass in "
                                  "-dynamic_dict during preprocess stage.")
 
-        return {
-            "output": output,
-            "target": batch.tgt[range_[0] + 1: range_[1], :, 0],
+        shard_state = super(CopyGeneratorLossCompute, self)._make_shard_state(
+            batch, output, 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, output, target, copy_attn, align,
+                      std_attn=None, coverage_attn=None):
         """Compute the loss.
 
         The args must match :func:`self._make_shard_state()`.
@@ -209,14 +214,18 @@ def _compute_loss(self, batch, output, target, copy_attn, align):
             copy_attn: the copy attention value.
             align: the align info.
         """
-
         target = target.view(-1)
         align = align.view(-1)
         scores = self.generator(
             self._bottle(output), self._bottle(copy_attn), batch.src_map
         )
         loss = self.criterion(scores, align, target)
 
+        if self.lambda_coverage != 0.0:
+            coverage_loss = self._compute_coverage_loss(std_attn,
+                                                        coverage_attn)
+            loss += coverage_loss
+
         # this block does not depend on the loss value computed above
         # and is used only for stats
         scores_data = collapse_copy_scores(

onmt/opts.py

@@ -176,8 +176,8 @@ def model_opts(parser):
               help="Divide copy loss by length of sequence")
     group.add('--coverage_attn', '-coverage_attn', action="store_true",
               help='Train a coverage attention layer.')
-    group.add('--lambda_coverage', '-lambda_coverage', type=float, default=1,
-              help='Lambda value for coverage.')
+    group.add('--lambda_coverage', '-lambda_coverage', type=float, default=0.0,
+              help='Lambda value for coverage loss of See et al (2017)')
     group.add('--loss_scale', '-loss_scale', type=float, default=0,
               help="For FP16 training, the static loss scale to use. If not "
                    "set, the loss scale is dynamically computed.")

onmt/utils/loss.py

@@ -25,6 +25,11 @@ def build_loss_compute(model, tgt_field, opt, train=True):
 
     padding_idx = tgt_field.vocab.stoi[tgt_field.pad_token]
     unk_idx = tgt_field.vocab.stoi[tgt_field.unk_token]
+
+    if opt.lambda_coverage != 0:
+        assert opt.coverage_attn, "--coverage_attn needs to be set in " \
+            "order to use --lambda_coverage != 0"
+
     if opt.copy_attn:
         criterion = onmt.modules.CopyGeneratorLoss(
             len(tgt_field.vocab), opt.copy_attn_force,
@@ -47,10 +52,12 @@ def build_loss_compute(model, tgt_field, opt, train=True):
     loss_gen = model.generator[0] if use_raw_logits else model.generator
     if opt.copy_attn:
         compute = onmt.modules.CopyGeneratorLossCompute(
-            criterion, loss_gen, tgt_field.vocab, opt.copy_loss_by_seqlength
+            criterion, loss_gen, tgt_field.vocab, opt.copy_loss_by_seqlength,
+            lambda_coverage=opt.lambda_coverage
         )
     else:
-        compute = NMTLossCompute(criterion, loss_gen)
+        compute = NMTLossCompute(
+            criterion, loss_gen, lambda_coverage=opt.lambda_coverage)
     compute.to(device)
 
     return compute
@@ -218,26 +225,53 @@ class NMTLossCompute(LossComputeBase):
     Standard NMT Loss Computation.
     """
 
-    def __init__(self, criterion, generator, normalization="sents"):
+    def __init__(self, criterion, generator, normalization="sents",
+                 lambda_coverage=0.0):
         super(NMTLossCompute, self).__init__(criterion, generator)
+        self.lambda_coverage = lambda_coverage
 
     def _make_shard_state(self, batch, output, range_, attns=None):
-        return {
+        shard_state = {
             "output": output,
             "target": batch.tgt[range_[0] + 1: range_[1], :, 0],
         }
+        if self.lambda_coverage != 0.0:
+            coverage = attns.get("coverage", None)
+            std = attns.get("std", None)
+            assert attns is not None
+            assert std is not None, "lambda_coverage != 0.0 requires " \
+                "attention mechanism"
+            assert coverage is not None, "lambda_coverage != 0.0 requires " \
+                "coverage attention"
+
+            shard_state.update({
+                "std_attn": attns.get("std"),
+                "coverage_attn": coverage
+            })
+        return shard_state
+
+    def _compute_loss(self, batch, output, target, std_attn=None,
+                      coverage_attn=None):
 
-    def _compute_loss(self, batch, output, target):
         bottled_output = self._bottle(output)
 
         scores = self.generator(bottled_output)
         gtruth = target.view(-1)
 
         loss = self.criterion(scores, gtruth)
+        if self.lambda_coverage != 0.0:
+            coverage_loss = self._compute_coverage_loss(
+                std_attn=std_attn, coverage_attn=coverage_attn)
+            loss += coverage_loss
         stats = self._stats(loss.clone(), scores, gtruth)
 
         return loss, stats
 
+    def _compute_coverage_loss(self, std_attn, coverage_attn):
+        covloss = torch.min(std_attn, coverage_attn).sum(2).view(-1)
+        covloss *= self.lambda_coverage
+        return covloss
+
 
 def filter_shard_state(state, shard_size=None):
     for k, v in state.items():