Bugfix: Loading via state_dict within lazy modules

test/nn/dense/test_linear.py

@@ -29,6 +29,24 @@ def test_lazy_linear(weight, bias):
     assert str(lin) == 'Linear(16, 32, bias=True)'
 
 
+@pytest.mark.parametrize('dim1,dim2', product([-1, 16], [-1, 16]))
+def test_load_lazy_linear(dim1, dim2):
+    lin1 = Linear(dim1, 32)
+    lin2 = Linear(dim1, 32)
+    lin2.load_state_dict(lin1.state_dict())
+
+    if dim1 != -1:
+        assert torch.allclose(lin1.weight, lin2.weight)
+        assert torch.allclose(lin1.bias, lin2.bias)
+        assert not hasattr(lin1, '_hook')
+        assert not hasattr(lin2, '_hook')
+    else:
+        assert isinstance(lin1.weight, UninitializedParameter)
+        assert isinstance(lin2.weight, UninitializedParameter)
+        assert hasattr(lin1, '_hook')
+        assert hasattr(lin2, '_hook')
+
+
 @pytest.mark.parametrize('lazy', [True, False])
 def test_identical_linear_default_initialization(lazy):
     x = torch.randn(3, 4, 16)

torch_geometric/nn/dense/linear.py

@@ -4,6 +4,7 @@
 import math
 
 import torch
+from torch import nn
 from torch import Tensor
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
@@ -49,7 +50,7 @@ def __init__(self, in_channels: int, out_channels: int, bias: bool = True,
         if in_channels > 0:
             self.weight = Parameter(torch.Tensor(out_channels, in_channels))
         else:
-            self.weight = torch.nn.parameter.UninitializedParameter()
+            self.weight = nn.parameter.UninitializedParameter()
             self._hook = self.register_forward_pre_hook(
                 self.initialize_parameters)
 
@@ -58,6 +59,9 @@ def __init__(self, in_channels: int, out_channels: int, bias: bool = True,
         else:
             self.register_parameter('bias', None)
 
+        self._load_hook = self._register_load_state_dict_pre_hook(
+            self._lazy_load_hook)
+
         self.reset_parameters()
 
     def __deepcopy__(self, memo):
@@ -71,45 +75,73 @@ def __deepcopy__(self, memo):
         return out
 
     def reset_parameters(self):
-        if self.in_channels > 0:
-            if self.weight_initializer == 'glorot':
-                inits.glorot(self.weight)
-            elif self.weight_initializer == 'uniform':
-                bound = 1.0 / math.sqrt(self.weight.size(-1))
-                torch.nn.init.uniform_(self.weight.data, -bound, bound)
-            elif self.weight_initializer == 'kaiming_uniform':
-                inits.kaiming_uniform(self.weight, fan=self.in_channels,
-                                      a=math.sqrt(5))
-            elif self.weight_initializer is None:
-                inits.kaiming_uniform(self.weight, fan=self.in_channels,
-                                      a=math.sqrt(5))
-            else:
-                raise RuntimeError(
-                    f"Linear layer weight initializer "
-                    f"'{self.weight_initializer}' is not supported")
-
-        if self.in_channels > 0 and self.bias is not None:
-            if self.bias_initializer == 'zeros':
-                inits.zeros(self.bias)
-            elif self.bias_initializer is None:
-                inits.uniform(self.in_channels, self.bias)
-            else:
-                raise RuntimeError(
-                    f"Linear layer bias initializer "
-                    f"'{self.bias_initializer}' is not supported")
+        if isinstance(self.weight, nn.parameter.UninitializedParameter):
+            pass
+        elif self.weight_initializer == 'glorot':
+            inits.glorot(self.weight)
+        elif self.weight_initializer == 'uniform':
+            bound = 1.0 / math.sqrt(self.weight.size(-1))
+            torch.nn.init.uniform_(self.weight.data, -bound, bound)
+        elif self.weight_initializer == 'kaiming_uniform':
+            inits.kaiming_uniform(self.weight, fan=self.in_channels,
+                                  a=math.sqrt(5))
+        elif self.weight_initializer is None:
+            inits.kaiming_uniform(self.weight, fan=self.in_channels,
+                                  a=math.sqrt(5))
+        else:
+            raise RuntimeError(f"Linear layer weight initializer "
+                               f"'{self.weight_initializer}' is not supported")
+
+        if isinstance(self.weight, nn.parameter.UninitializedParameter):
+            pass
+        elif self.bias is None:
+            pass
+        elif self.bias_initializer == 'zeros':
+            inits.zeros(self.bias)
+        elif self.bias_initializer is None:
+            inits.uniform(self.in_channels, self.bias)
+        else:
+            raise RuntimeError(f"Linear layer bias initializer "
+                               f"'{self.bias_initializer}' is not supported")
 
     def forward(self, x: Tensor) -> Tensor:
         """"""
         return F.linear(x, self.weight, self.bias)
 
     @torch.no_grad()
     def initialize_parameters(self, module, input):
-        if isinstance(self.weight, torch.nn.parameter.UninitializedParameter):
+        if isinstance(self.weight, nn.parameter.UninitializedParameter):
             self.in_channels = input[0].size(-1)
             self.weight.materialize((self.out_channels, self.in_channels))
             self.reset_parameters()
-        module._hook.remove()
-        delattr(module, '_hook')
+        self._hook.remove()
+        delattr(self, '_hook')
+
+    def _save_to_state_dict(self, destination, prefix, keep_vars):
+        if isinstance(self.weight, nn.parameter.UninitializedParameter):
+            destination[prefix + 'weight'] = self.weight
+        else:
+            destination[prefix + 'weight'] = self.weight.detach()
+        if self.bias is not None:
+            destination[prefix + 'bias'] = self.bias.detach()
+
+    def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,
+                        missing_keys, unexpected_keys, error_msgs):
+
+        weight = state_dict[prefix + 'weight']
+        if isinstance(weight, nn.parameter.UninitializedParameter):
+            self.in_channels = -1
+            self.weight = nn.parameter.UninitializedParameter()
+            if not hasattr(self, '_hook'):
+                self._hook = self.register_forward_pre_hook(
+                    self.initialize_parameters)
+
+        elif isinstance(self.weight, nn.parameter.UninitializedParameter):
+            self.in_channels = weight.size(-1)
+            self.weight.materialize((self.out_channels, self.in_channels))
+            if hasattr(self, '_hook'):
+                self._hook.remove()
+                delattr(self, '_hook')
 
     def __repr__(self) -> str:
         return (f'{self.__class__.__name__}({self.in_channels}, '