Repository: pytorch/pytorch

Issue:
[bug?] Problem with load_state_dict after installing latest pytorch by source
- OS: MacOS Sierra
- PyTorch version: 0.4.0a0+7588893
- How you installed PyTorch (conda, pip, source):  source
- Python version: Python 2.7.14 |Anaconda, Inc.|
- CUDA/cuDNN version: No Cuda
- GPU models and configuration: No GPU models
- GCC version (if compiling from source): GCC 4.2.1 Compatible Clang 4.0.1

In order to use onnx on nn.InstanceNorm2d layers, I installed the latest pytorch by source following the instruction.
Before installing the latest pytorch from source.  I trained some pytorch model using the pytorch version installed through http://pytorch.org/. I was able to save the trained model as .pth file and load it back in successfully by doing the following:

`model = myNet(parameters...)
model.load_state_dict(torch.load('xxx.pth'))`

However, after I install the latest pytorch (from source). The load_state_dict is reporting errors:

> File "/anaconda2/lib/python2.7/site-packages/torch/nn/modules/module.py", line 528, in load_state_dict
    .format(name))
KeyError: 'unexpected key "pre_model.1.running_mean" in state_dict'

I did the following to check which layer was causing the problem
`model1 = myNet(parameters...)`
`model2 = torch.load('xxx.pth')`

`for name, param in model2.items():`
`     print(name)`

`for name, param in model1.state_dict().items():`
`      print(name)`

I found that for the instanceNorm2D layers, model2 contains xxx.running_mean and xxx.running_var, but not model1. Any idea what I should do? 





Discussion:
The `running_*` are disabled by default for InstanceNorm* layers after #4922 . You can add `track_running_stats=True` to the InstanceNorm* layer constructors.
that works. Thank you.

Resolving patch:
diff --git a/torch/nn/modules/instancenorm.py b/torch/nn/modules/instancenorm.py
--- a/torch/nn/modules/instancenorm.py
+++ b/torch/nn/modules/instancenorm.py
@@ -11,6 +11,23 @@ def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=False,
     def _check_input_dim(self, input):
         return NotImplemented
 
+    def _load_state_dict_key_mismatch(self, full_name, name, is_missing):
+        if not is_missing and not self.track_running_stats and \
+                name in ('running_mean', 'running_var'):
+            raise KeyError(
+                'Unexpected running stats buffer "{name}" in state_dict for '
+                '{klass} with track_running_stats=False. If you are trying to '
+                'load a checkpoint saved before 0.4.0, this may be expected '
+                'because {klass} does not track running stats by default '
+                'anymore since 0.4.0. If running stats are not used, remove '
+                'them from state_dict before calling load_state_dict. '
+                'Otherwise, set track_running_stats=True in {klass} to use '
+                'running stats. See the documentation of {klass} for more '
+                'details.'
+                .format(name=full_name, klass=self.__class__.__name__))
+        super(_InstanceNorm, self)._load_state_dict_key_mismatch(
+            full_name, name, is_missing)
+
     def forward(self, input):
         self._check_input_dim(input)
 
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
--- a/torch/nn/modules/module.py
+++ b/torch/nn/modules/module.py
@@ -507,6 +507,20 @@ def state_dict(self, destination=None, prefix='', keep_vars=False):
                 module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
         return destination
 
+    def _load_state_dict_key_mismatch(self, full_name, name, is_missing):
+        r"""This is called in :meth:`~torch.nn.Module.load_state_dict` when
+        there is state dict key mismatch in ``strict=True`` mode. This method
+        can be overridden by subclasses to raise class-specific errors.
+
+        When :attr:`is_missing` is ``True``, :attr:`full_name` can not be found in
+        the dict being loaded. When :attr:`is_missing` is ``False``,
+        :attr:`full_name` is unexpected in the dict being loaded.
+
+        :attr:`name` is the actual name of the parameter/buffer, i.e., the
+        substring after the last `dot` in :attr:`full_name`.
+        """
+        pass
+
     def load_state_dict(self, state_dict, strict=True):
         r"""Copies parameters and buffers from :attr:`state_dict` into
         this module and its descendants. If :attr:`strict` is ``True`` then
@@ -520,6 +534,17 @@ def load_state_dict(self, state_dict, strict=True):
                 match the keys returned by this module's `:func:`state_dict()`
                 function.
         """
+        def submodule_key_mismatch(full_name, is_missing):
+            module = self
+            names = full_name.split(".")
+            for module_name in names[:-1]:
+                if module_name in module._modules:
+                    module = module._modules[module_name]
+                else:
+                    return
+            module._load_state_dict_key_mismatch(full_name, names[-1], is_missing)
+
+        unexpected = []
         own_state = self.state_dict()
         for name, param in state_dict.items():
             if name in own_state:
@@ -534,12 +559,24 @@ def load_state_dict(self, state_dict, strict=True):
                                        'whose dimensions in the checkpoint are {}.'
                                        .format(name, own_state[name].size(), param.size()))
             elif strict:
-                raise KeyError('unexpected key "{}" in state_dict'
-                               .format(name))
+                unexpected.append(name)
         if strict:
             missing = set(own_state.keys()) - set(state_dict.keys())
+            # pass the mismatch info to submodules so that they have a chance to
+            # raise a custom class-specific error
+            for name in unexpected:
+                submodule_key_mismatch(name, False)
+            for name in missing:
+                submodule_key_mismatch(name, True)
+            error_msg = ''
+            if len(unexpected) > 0:
+                error_msg += 'Unexpected key(s) in state_dict: {}. '.format(
+                    ', '.join('"{}"'.format(k) for k in unexpected))
             if len(missing) > 0:
-                raise KeyError('missing keys in state_dict: "{}"'.format(missing))
+                error_msg += 'Missing key(s) in state_dict: {}. '.format(
+                    ', '.join('"{}"'.format(k) for k in unexpected))
+            if len(error_msg) > 0:
+                raise KeyError(error_msg)
 
     def parameters(self):
         r"""Returns an iterator over module parameters.


Test patch: