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".
```


Discussion:
We may have some compatibility issues.
Could you test it with v.0.6.0?

> We may have some compatibility issues.
> Could you test it with v.0.6.0?

Thank you for your reply. I have test it on  v.0.6.0, but the same issue still occurs. By the way, I read the RESULTS.md file and the pretained model (pytorch Transformer, accum grad 8, single GPU) I use is based on v.0.3.1. However, run.sh in v.0.3.1 doesn't import the training and decoding config files, which are listed in the  RESULTS.md file. So, should I try to  test it with v.0.3.1 or there is other way to run it correctly? 
>  I use is based on v.0.3.1

The name was changed in v.0.4.0 https://github.com/espnet/espnet/commit/21d70286c354c66c0350e65dc098d2ee236faccc#diff-bffb1396f038b317b2b64dd96e6d3563

A possible fix is renaming the state dict and upload it to the URL. Or we can rename the key on the fly like what we have done for positional encoding before https://github.com/espnet/espnet/blob/a97a270aa581d7788bf4d71d108a1e5bd3d23f7d/espnet/nets/pytorch_backend/transformer/embedding.py#L45

Resolving patch:
diff --git a/espnet/nets/pytorch_backend/nets_utils.py b/espnet/nets/pytorch_backend/nets_utils.py
--- 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
--- 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.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 @@ def __init__(self, odim,
                  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
--- 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.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 @@ def __init__(self, idim,
                  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(


Test patch:
diff --git a/test/test_e2e_asr_transformer.py b/test/test_e2e_asr_transformer.py
--- a/test/test_e2e_asr_transformer.py
+++ b/test/test_e2e_asr_transformer.py
@@ -6,6 +6,8 @@
 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()