return encoder representations only if necessary

OpenNMT · Feb 17, 2020 · ff858ac · ff858ac
1 parent 306b2e5
commit ff858ac
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 19 deletions.
diff --git a/onmt/models/model.py b/onmt/models/model.py
@@ -46,22 +46,23 @@ def forward(self, src, tgt, lengths, bptt=False,
 
         enc_state, memory_bank, lengths = self.encoder(src, lengths)
 
+        if bptt is False:
+            self.decoder.init_state(src, memory_bank, enc_state)
+
+        dec_out, attns = self.decoder(dec_in, memory_bank,
+                                      memory_lengths=lengths,
+                                      with_align=with_align)
+
         if encode_tgt:
             # tgt for zero shot alignment loss
             tgt_lengths = torch.Tensor(tgt.size(1))\
                                .type_as(memory_bank) \
                                .long() \
                                .fill_(tgt.size(0))
             embs_tgt, memory_bank_tgt, ltgt = self.encoder(tgt, tgt_lengths)
-        else:
-            memory_bank_tgt = None
+            return dec_out, attns, memory_bank, memory_bank_tgt
 
-        if bptt is False:
-            self.decoder.init_state(src, memory_bank, enc_state)
-        dec_out, attns = self.decoder(dec_in, memory_bank,
-                                      memory_lengths=lengths,
-                                      with_align=with_align)
-        return dec_out, attns, memory_bank, memory_bank_tgt
+        return dec_out, attns
 
     def update_dropout(self, dropout):
         self.encoder.update_dropout(dropout)

diff --git a/onmt/trainer.py b/onmt/trainer.py
@@ -317,13 +317,21 @@ def validate(self, valid_iter, moving_average=None):
                 tgt = batch.tgt
 
                 # F-prop through the model.
-                outputs, attns, enc_src, enc_tgt = valid_model(
-                    src, tgt, src_lengths,
-                    with_align=self.with_align)
+                if self.encode_tgt:
+                    outputs, attns, enc_src, enc_tgt = valid_model(
+                        src, tgt, src_lengths,
+                        with_align=self.with_align,
+                        encode_tgt=self.encode_tgt)
+                else:
+                    output, attns = valid_model(
+                        src, tgt, src_lengths,
+                        with_align=self.with_align)
+                    enc_src, enc_tgt = None, None
 
                 # Compute loss.
                 _, batch_stats = self.valid_loss(
-                    batch, outputs, attns, enc_src, enc_tgt)
+                    batch, outputs, attns,
+                    enc_src=enc_src, enc_tgt=enc_tgt)
 
                 # Update statistics.
                 stats.update(batch_stats)
@@ -366,9 +374,16 @@ def _gradient_accumulation(self, true_batches, normalization, total_stats,
                 if self.accum_count == 1:
                     self.optim.zero_grad()
 
-                outputs, attns, enc_src, enc_tgt = self.model(
-                    src, tgt, src_lengths, bptt=bptt,
-                    with_align=self.with_align, encode_tgt=self.encode_tgt)
+                is self.encode_tgt:
+                    outputs, attns, enc_src, enc_tgt = self.model(
+                        src, tgt, src_lengths, bptt=bptt,
+                        with_align=self.with_align, encode_tgt=self.encode_tgt)
+                else:
+                    output, attns = self.model(
+                        src, tgt, src_lengths, bptt=bptt,
+                        with_align=self.with_align)
+                    enc_src, enc_tgt = None, None
+
                 bptt = True
 
                 # 3. Compute loss.
@@ -377,8 +392,8 @@ def _gradient_accumulation(self, true_batches, normalization, total_stats,
                         batch,
                         outputs,
                         attns,
-                        enc_src,
-                        enc_tgt,
+                        enc_src=enc_src,
+                        enc_tgt=enc_tgt,
                         normalization=normalization,
                         shard_size=self.shard_size,
                         trunc_start=j,

diff --git a/onmt/utils/loss.py b/onmt/utils/loss.py
@@ -124,8 +124,8 @@ def __call__(self,
                  batch,
                  output,
                  attns,
-                 enc_src,
-                 enc_tgt,
+                 enc_src=None,
+                 enc_tgt=None,
                  normalization=1.0,
                  shard_size=0,
                  trunc_start=0,