Repository: pfnet/pytorch-pfn-extras

Issue:
`autoload=True` not working with models in the gpu
If we create the snapshot extension with `autoload=True` the model will not correctly load its state.

`autoload=true` loads the state in the CPU. but it is not executed until the first iteration and it will overwrite the device of the model.  This requires us to call `start_extensions` manually and then do the device move for it to work

```python
# model parameters will be moved to cpu at the beginning of the first iteration due autoload being executed there
...
manager.extend(extensions.snapshot(autoload=True)
# move the model, but this will be overwriten later
model.cuda()

for batch in train_loader:
    with manager.run_iteration():   # Snapshot load happens the first time this is executed
        model(batch.cuda())   #  Error! weights are in the cpu again
```


Discussion:
Actually this only happens with ppe.Lazy modules.

Resolving patch:
diff --git a/pytorch_pfn_extras/nn/modules/lazy.py b/pytorch_pfn_extras/nn/modules/lazy.py
--- a/pytorch_pfn_extras/nn/modules/lazy.py
+++ b/pytorch_pfn_extras/nn/modules/lazy.py
@@ -81,32 +81,45 @@ def _lazy_load_hook(
             missing_keys, unexpected_keys, error_msgs):
         """load_state_dict pre-hook function for lazy buffers and parameters.
 
-        The purpose of this hook is to adjust the current state and/or
-        ``state_dict`` being loaded so that a module instance serialized in
-        both un/initialized state can be deserialized onto both un/initialized
-        module instance.
+        The purpose of this hook is to check the current state and/or
+        ``state_dict`` being loaded and ensure that both are states
+        are properly initialized.
 
         See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``
         for the details of the hook specification.
         """
         for name in self.lazy_buffer_names:
-            # Avoid shape mismatch error when loading an initialized buffer
-            # onto an uninitialized module instance.
-            self.register_buffer(name, state_dict[prefix + name])
+            key = prefix + name
+            module_initialized = getattr(self, name).shape != (0,)
+            state_initialized = state_dict[key].shape != (0,)
+            if module_initialized and not state_initialized:
+                raise RuntimeError(
+                    'Can\'t load non-initialized buffers in already '
+                    'initialized modules')
+            elif not module_initialized and state_initialized:
+                # Here we need to avoid a tensor size mismatch
+                # this is a regular tensor without a materialize
+                # method, so we can just resize for the load logic to copy
+                # the contents later to the correct device the module
+                # was moved to
+                getattr(self, name).resize_(state_dict[key].size())
 
         for name in self.lazy_parameter_names:
-            # The parameter may not exist in the loaded ``state_dict`` if the
+            # The parameter does not exist in the loaded ``state_dict`` if the
             # original module was serialized before initializing lazy
             # parameters (see comments of ``state_dict``).
             key = prefix + name
-            if key in state_dict:
-                # The model was serialized after initialization.
-                self.register_parameter(
-                    name, torch.nn.Parameter(state_dict[key]))
-            else:
-                # The model was serialized before initialization.
+            module_initialized = not isinstance(
+                getattr(self, name), UninitializedParameter)
+            state_initialized = key in state_dict
+            if module_initialized and not state_initialized:
+                raise RuntimeError(
+                    'Can\'t load uninitialized parameters in already '
+                    'initialized modules')
+            elif not module_initialized and state_initialized:
+                getattr(self, name).materialize(state_dict[key].shape)
+            elif key not in state_dict and not module_initialized:
                 param = UninitializedParameter()
-                self.register_parameter(name, param)
                 state_dict[key] = param
 
 
@@ -133,3 +146,25 @@ def is_leaf(self):
     Use of uninitialized lazy parameter in Optimizer has been detected.
     Maybe you forgot to run forward before passing `module.parameters()` to the optimizer?''')  # NOQA
         return True
+
+    def materialize(self, shape, device=None, dtype=None):
+        r"""Create a Parameter with the same properties of the uninitialized
+        one. Given a shape, it materializes a parameter in the same device
+        and with the same `dtype` as the current one or the specified ones in
+        the arguments.
+
+        Args:
+            shape : (tuple): the shape for the materialized tensor.
+            device (:class:`torch.device`): the desired device of the
+                parameters
+                and buffers in this module. Optional.
+            dtype (:class:`torch.dtype`): the desired floating point type of
+                the floating point parameters and buffers in this module.
+                Optional.
+        """
+        if device is None:
+            device = self.data.device
+        if dtype is None:
+            dtype = self.data.dtype
+        self.data = torch.empty(shape, device=device, dtype=dtype)
+        self.__class__ = torch.nn.Parameter


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy.py b/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy.py
--- a/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy.py
+++ b/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy.py
@@ -136,17 +136,41 @@ def test_save_load(self, init_src, init_dst):
         input = self.get_input()
         model_src = self.get_lazy_module()
         model_dst = self.get_lazy_module()
-
         if init_src:
             torch.manual_seed(0)
             model_src(input)
+
         if init_dst:
             torch.manual_seed(0)
             model_dst(input)
+            module_params = [
+                getattr(model_dst, name)
+                for name in model_dst.lazy_parameter_names
+            ]
+            module_buffers = [
+                getattr(model_dst, name)
+                for name in model_dst.lazy_buffer_names
+            ]
 
         with tempfile.NamedTemporaryFile(delete=False) as f:
             torch.save(model_src.state_dict(), f.name)
-            model_dst.load_state_dict(torch.load(f.name))
+            if not init_src and init_dst:
+                with pytest.raises(RuntimeError):
+                    model_dst.load_state_dict(torch.load(f.name))
+            else:
+                model_dst.load_state_dict(torch.load(f.name))
+
+        # Ensure that if the model was initialized, the parameters are the same
+        # after loading a state dict
+        if init_dst:
+            for name, param in zip(
+                model_dst.lazy_parameter_names, module_params
+            ):
+                assert getattr(model_dst, name).data_ptr() == param.data_ptr()
+            for name, buffer in zip(
+                model_dst.lazy_buffer_names, module_buffers
+            ):
+                assert getattr(model_dst, name).data_ptr() == buffer.data_ptr()
 
         torch.manual_seed(0)
         expected = model_src(input)