Repository: espnet/espnet
Issue: Missing and unexpected key error in decoding

Hi all,
I'm the beginner for ESpnet and I downloaded librispeech pretrained model and run the stage 5 (decoding) of run.sh. Then, the following error occurs. It seems that this issue is similar to  issue #154, but I'm not so sure about it. 

```
Traceback (most recent call last):
  File "/home/training/r08943020/espnet/egs/librispeech/asr1/../../../espnet/bin/asr_recog.py", line 208, in <module>
    main(sys.argv[1:])
  File "/home/training/r08943020/espnet/egs/librispeech/asr1/../../../espnet/bin/asr_recog.py", line 185, in main
    recog_v2(args)
  File "/home/training/r08943020/espnet/espnet/asr/pytorch_backend/recog.py", line 39, in recog_v2
    model, train_args = load_trained_model(args.model)
  File "/home/training/r08943020/espnet/espnet/asr/pytorch_backend/asr_init.py", line 140, in load_trained_model
    torch_load(model_path, model)
  File "/home/training/r08943020/espnet/espnet/asr/asr_utils.py", line 528, in torch_load
    model.load_state_dict(remove_dataparallel(model_state_dict))
  File "/home/training/r08943020/espnet/tools/venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 777, in load_state_dict
    self.__class__.__name__, "\\n\\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for E2E:
        Missing key(s) in state_dict: "encoder.embed.conv.0.weight", "encoder.embed.conv.0.bias", "encoder.embed.conv.2.weight", "encoder.embed.conv.2.bias", "encoder.embed.out.0.weight", "encoder.embed.out.0.bias", "encoder.after_norm.weight", "encoder.after_norm.bias", "decoder.after_norm.weight", "decoder.after_norm.bias".
        Unexpected key(s) in state_dict: "encoder.input_layer.conv.0.weight", "encoder.input_layer.conv.0.bias", "encoder.input_layer.conv.2.weight", "encoder.input_layer.conv.2.bias", "encoder.input_layer.out.0.weight", "encoder.input_layer.out.0.bias", "encoder.input_layer.out.1.pe", "encoder.norm.weight", "encoder.norm.bias", "decoder.output_norm.weight", "decoder.output_norm.bias".
```

Changed files:
test/test_e2e_asr_transformer.py, espnet/nets/pytorch_backend/nets_utils.py, espnet/nets/pytorch_backend/transformer/decoder.py, espnet/nets/pytorch_backend/transformer/encoder.py

diff --git a/espnet/nets/pytorch_backend/nets_utils.py b/espnet/nets/pytorch_backend/nets_utils.py
index ca9c5d65b..098e3b4c5 100644
--- a/espnet/nets/pytorch_backend/nets_utils.py
+++ b/espnet/nets/pytorch_backend/nets_utils.py
@@ -3,6 +3,8 @@
 """Network related utility tools."""
 
 import logging
+from typing import Dict
+
 import numpy as np
 import torch
 
@@ -435,3 +437,15 @@ def get_subsample(train_args, mode, arch):
 
     else:
         raise ValueError('Invalid options: mode={}, arch={}'.format(mode, arch))
+
+
+def rename_state_dict(old_prefix: str, new_prefix: str, state_dict: Dict[str, torch.Tensor]):
+    """Replace keys of old prefix with new prefix in state dict."""
+    # need this list not to break the dict iterator
+    old_keys = [k for k in state_dict if k.startswith(old_prefix)]
+    if len(old_keys) > 0:
+        logging.warning(f'Rename: {old_prefix} -> {new_prefix}')
+    for k in old_keys:
+        v = state_dict.pop(k)
+        new_k = k.replace(old_prefix, new_prefix)
+        state_dict[new_k] = v
diff --git a/espnet/nets/pytorch_backend/transformer/decoder.py b/espnet/nets/pytorch_backend/transformer/decoder.py
index ae62e6f1e..47d84680e 100644
--- a/espnet/nets/pytorch_backend/transformer/decoder.py
+++ b/espnet/nets/pytorch_backend/transformer/decoder.py
@@ -8,6 +8,7 @@
 
 import torch
 
+from espnet.nets.pytorch_backend.nets_utils import rename_state_dict
 from espnet.nets.pytorch_backend.transformer.attention import MultiHeadedAttention
 from espnet.nets.pytorch_backend.transformer.decoder_layer import DecoderLayer
 from espnet.nets.pytorch_backend.transformer.embedding import PositionalEncoding
@@ -18,6 +19,12 @@ from espnet.nets.pytorch_backend.transformer.repeat import repeat
 from espnet.nets.scorer_interface import ScorerInterface
 
 
+def _pre_hook(state_dict, prefix, local_metadata, strict,
+              missing_keys, unexpected_keys, error_msgs):
+    # https://github.com/espnet/espnet/commit/3d422f6de8d4f03673b89e1caef698745ec749ea#diff-bffb1396f038b317b2b64dd96e6d3563
+    rename_state_dict(prefix + "output_norm.", prefix + "after_norm.", state_dict)
+
+
 class Decoder(ScorerInterface, torch.nn.Module):
     """Transfomer decoder module.
 
@@ -53,6 +60,7 @@ class Decoder(ScorerInterface, torch.nn.Module):
                  concat_after=False):
         """Construct an Decoder object."""
         torch.nn.Module.__init__(self)
+        self._register_load_state_dict_pre_hook(_pre_hook)
         if input_layer == "embed":
             self.embed = torch.nn.Sequential(
                 torch.nn.Embedding(odim, attention_dim),
diff --git a/espnet/nets/pytorch_backend/transformer/encoder.py b/espnet/nets/pytorch_backend/transformer/encoder.py
index b89382e2d..635070c61 100644
--- a/espnet/nets/pytorch_backend/transformer/encoder.py
+++ b/espnet/nets/pytorch_backend/transformer/encoder.py
@@ -8,8 +8,8 @@
 
 import torch
 
+from espnet.nets.pytorch_backend.nets_utils import rename_state_dict
 from espnet.nets.pytorch_backend.transducer.vgg import VGG2L
-
 from espnet.nets.pytorch_backend.transformer.attention import MultiHeadedAttention
 from espnet.nets.pytorch_backend.transformer.embedding import PositionalEncoding
 from espnet.nets.pytorch_backend.transformer.encoder_layer import EncoderLayer
@@ -21,6 +21,14 @@ from espnet.nets.pytorch_backend.transformer.repeat import repeat
 from espnet.nets.pytorch_backend.transformer.subsampling import Conv2dSubsampling
 
 
+def _pre_hook(state_dict, prefix, local_metadata, strict,
+              missing_keys, unexpected_keys, error_msgs):
+    # https://github.com/espnet/espnet/commit/21d70286c354c66c0350e65dc098d2ee236faccc#diff-bffb1396f038b317b2b64dd96e6d3563
+    rename_state_dict(prefix + "input_layer.", prefix + "embed.", state_dict)
+    # https://github.com/espnet/espnet/commit/3d422f6de8d4f03673b89e1caef698745ec749ea#diff-bffb1396f038b317b2b64dd96e6d3563
+    rename_state_dict(prefix + "norm.", prefix + "after_norm.", state_dict)
+
+
 class Encoder(torch.nn.Module):
     """Transformer encoder module.
 
@@ -60,6 +68,7 @@ class Encoder(torch.nn.Module):
                  padding_idx=-1):
         """Construct an Encoder object."""
         super(Encoder, self).__init__()
+        self._register_load_state_dict_pre_hook(_pre_hook)
 
         if input_layer == "linear":
             self.embed = torch.nn.Sequential(
diff --git a/test/test_e2e_asr_transformer.py b/test/test_e2e_asr_transformer.py
index 67f9345ad..763e3a595 100644
--- a/test/test_e2e_asr_transformer.py
+++ b/test/test_e2e_asr_transformer.py
@@ -6,6 +6,8 @@ import numpy
 import pytest
 import torch
 
+from espnet.nets.pytorch_backend.nets_utils import rename_state_dict
+
 
 logging.basicConfig(
     level=logging.INFO, format='%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s')
@@ -304,5 +306,19 @@ def test_transformer_parallel():
     assert max_acc > 0.8
 
 
+# https://github.com/espnet/espnet/issues/1750
+def test_v0_3_transformer_input_compatibility():
+    args = make_arg()
+    model, x, ilens, y, data = prepare("pytorch", args)
+    # these old names are used in v.0.3.x
+    state_dict = model.state_dict()
+    prefix = "encoder."
+    rename_state_dict(prefix + "embed.", prefix + "input_layer.", state_dict)
+    rename_state_dict(prefix + "after_norm.", prefix + "norm.", state_dict)
+    prefix = "decoder."
+    rename_state_dict(prefix + "after_norm.", prefix + "output_norm.", state_dict)
+    model.load_state_dict(state_dict)
+
+
 if __name__ == "__main__":
     run_transformer_copy()