Kurz vorm durchdrehen

lib/modules/utils.py

@@ -67,6 +67,23 @@ class AutoPad(nn.Module):
         return x
 
 
+class WeightInit:
+
+    def __init__(self, in_place_init_function):
+        self.in_place_init_function = in_place_init_function
+
+    def __call__(self, m):
+        if hasattr(m, 'weight'):
+            if isinstance(m.weight, torch.Tensor):
+                if m.weight.ndim < 2:
+                    m.weight.data.fill_(0.01)
+                else:
+                    self.in_place_init_function(m.weight)
+        if hasattr(m, 'bias'):
+            if isinstance(m.bias, torch.Tensor):
+                m.bias.data.fill_(0.01)
+
+
 class LightningBaseModule(pl.LightningModule, ABC):
 
     @classmethod
@@ -128,15 +145,9 @@ class LightningBaseModule(pl.LightningModule, ABC):
     def test_epoch_end(self, outputs):
         raise NotImplementedError
 
-    def init_weights(self):
-        def _weight_init(m):
-            if hasattr(m, 'weight'):
-                if isinstance(m.weight, torch.Tensor):
-                    torch.nn.init.xavier_uniform_(m.weight)
-            if hasattr(m, 'bias'):
-                if isinstance(m.bias, torch.Tensor):
-                    m.bias.data.fill_(0.01)
-        self.apply(_weight_init)
+    def init_weights(self, in_place_init_func_=nn.init.xavier_uniform_):
+        weight_initializer = WeightInit(in_place_init_function=in_place_init_func_)
+        self.apply(weight_initializer)
 
     # Dataloaders
     # ================================================================================