Repository: OpenNMT/OpenNMT-py

Issue:
TransformerDecoder class layer cache
In v3.0 in the TransformerDecoder class, the layer cache is no longer reset. This leads to an error when switiching between training and inference. During inference, in the layer cache the keys, etc. of the last tokens are saved, but they need to be removed when switching back to training.


Discussion:
Hi @MaxThFe ,
Can you please be more specific in what context (training Inference training) you are seeing this ?
@vince62s Sorry about being unclear. This happens when inferencing first and then training the model. I believe this is due to the following change:

In version 2.3.0, in onmt.decoders.TransformerDecoder.forward, the forward is implemented with a cache check:
for i, layer in enumerate(self.transformer_layers):
            layer_cache = (
                self.state["cache"]["layer_{}".format(i)]
                if step is not None
                else None
            )
            output, attn, attn_align = layer(
                output,
                src_memory_bank,
                src_pad_mask,
                tgt_pad_mask,
                layer_cache=layer_cache,
                step=step,
                with_align=with_align,
            )

In version 3.0 however this part is changed:
        for layer in self.transformer_layers:
            dec_out, attn, attn_align = layer(
                dec_out,
                enc_out,
                src_pad_mask,
                tgt_pad_mask,
                step=step,
                with_align=with_align,
            )
            if attn_align is not None:
                attn_aligns.append(attn_align)

Thus, when calling the forward function during inference, the cache of each layer of self attention is filled. Afterwards however, when calling the forward function in training, this leads to an error as there is layer cache. This is unwanted behaviour, as denoted in onmt.modules.MultiHeadedAttention.forward as well:
# 1) Project key, value, and query.
        # as a reminder at training layer_cache[0] remains False
        if self.layer_cache[0]:

Hope this is sufficiently clear. 
I know but in what context are you experiencing this ?
for instance, when we train with the new scroing stuff at validation time, then we do this:
https://github.com/OpenNMT/OpenNMT-py/blob/master/onmt/utils/scoring_utils.py#L163-L172
which indeed reset the cache for training after having inferred for scoring.
@vince62s oh I see what you mean. I used only the decoder model in a custom training framework.
But overall you are correct, there is a slight change in behavior if we consider the API as a library for custom frameworks.
Can you PR ?
I think the easiest would be a test on step == None to reset the cache.

Resolving patch:
diff --git a/onmt/decoders/transformer.py b/onmt/decoders/transformer.py
--- a/onmt/decoders/transformer.py
+++ b/onmt/decoders/transformer.py
@@ -440,6 +440,14 @@ def forward(self, tgt, enc_out=None, step=None, **kwargs):
             enc_out = self.embeddings(tgt)
         if step == 0:
             self._init_cache(enc_out)
+        elif step is None:
+            for layer in self.transformer_layers:
+                layer.self_attn.layer_cache = (
+                    False, {'keys': torch.tensor([]),
+                            'values': torch.tensor([])})
+                layer.context_attn.layer_cache = (
+                    False, {'keys': torch.tensor([]),
+                            'values': torch.tensor([])})
 
         tgt_words = tgt[:, :, 0]
 
@@ -627,6 +635,11 @@ def forward(self, tgt, enc_out=None, step=None, **kwargs):
         """Decode, possibly stepwise."""
         if step == 0:
             self._init_cache(tgt)
+        elif step is None:
+            for layer in self.transformer_layers:
+                layer.self_attn.layer_cache = (
+                    False, {'keys': torch.tensor([]),
+                            'values': torch.tensor([])})
 
         tgt_words = tgt[:, :, 0]
 
diff --git a/onmt/utils/scoring_utils.py b/onmt/utils/scoring_utils.py
--- a/onmt/utils/scoring_utils.py
+++ b/onmt/utils/scoring_utils.py
@@ -1,6 +1,5 @@
 import codecs
 import os
-import torch
 from onmt.utils.parse import ArgumentParser
 from onmt.translate import GNMTGlobalScorer, Translator
 from onmt.opts import translate_opts
@@ -160,14 +159,4 @@ def translate(self, model, transformed_batches, gpu_rank, step, mode):
                     file.write("SOURCE: {}\n".format(raw_sources[i]))
                     file.write("REF: {}\n".format(texts_ref[i]))
                     file.write("PRED: {}\n\n".format(preds[i]))
-        # We deactivate the decoder's cache
-        # as we use teacher forcing at training time.
-        if hasattr(model.decoder, 'transformer_layers'):
-            for layer in model.decoder.transformer_layers:
-                layer.self_attn.layer_cache = (False,
-                                               {'keys': torch.tensor([]),
-                                                'values': torch.tensor([])})
-                layer.context_attn.layer_cache = (False,
-                                                  {'keys': torch.tensor([]),
-                                                   'values': torch.tensor([])})
         return preds, texts_ref


Test patch: