shifter init unit test

Author: NathanSoeding

neuralpredictors/layers/shifters/mlp.py

@@ -3,7 +3,7 @@
 import torch
 from torch import nn
 from torch.nn import ModuleDict
-from torch.nn.init import xavier_normal
+from torch.nn.init import xavier_normal_
 
 from .base import Shifter
 
@@ -37,11 +37,12 @@ def regularizer(self):
         return 0
 
     def initialize(self):
-        for layer in self.mlp:
-            if isinstance(layer, nn.Linear):
-                xavier_normal_(layer.weight)
-                if layer.bias is not None:
-                    nn.init.zeros_(layer.bias)
+        with torch.no_grad():
+            for layer in self.mlp:
+                if isinstance(layer, nn.Linear):
+                    xavier_normal_(layer.weight)
+                    if layer.bias is not None:
+                        nn.init.zeros_(layer.bias)
 
     def forward(self, pupil_center, trial_idx=None):
         if trial_idx is not None:
@@ -77,8 +78,5 @@ def __init__(
         for k in data_keys:
             self.add_module(k, MLP(input_channels, hidden_channels_shifter, shift_layers, bias))
 
-    def initialize(self, **kwargs):
-        pass
-
     def regularizer(self, data_key):
         return self[data_key].regularizer() * self.gamma_shifter

test/test_shifter_init.py

@@ -0,0 +1,12 @@
+import torch
+from torch import nn
+from unittest.mock import patch
+
+from neuralpredictors.layers.shifters import MLP
+
+def test_xavier_initialization_is_used():
+    with patch("neuralpredictors.layers.shifters.mlp.xavier_normal_") as mock_xavier:
+        MLP(shift_layers=3)
+
+    # One Linear layer per shift layer
+    assert mock_xavier.call_count == 3