Allow to set tensors parameter to control what is initialized (#104)

deeplay/initializers/constant.py

@@ -21,8 +21,5 @@ def __init__(
         self.weight = weight
         self.bias = bias
 
-    def initialize_weight(self, tensor):
+    def initialize_tensor(self, tensor, name):
         tensor.data.fill_(self.weight)
-
-    def initialize_bias(self, tensor):
-        tensor.data.fill_(self.bias)

deeplay/initializers/initializer.py

@@ -3,15 +3,11 @@ class Initializer:
     def __init__(self, targets):
         self.targets = targets
 
-    def initialize(self, module):
+    def initialize(self, module, tensors=("weight", "bias")):
         if isinstance(module, self.targets):
-            if hasattr(module, "weight") and module.weight is not None:
-                self.initialize_weight(module.weight)
-            if hasattr(module, "bias") and module.bias is not None:
-                self.initialize_bias(module.bias)
+            for tensor in tensors:
+                if hasattr(module, tensor) and getattr(module, tensor) is not None:
+                    self.initialize_tensor(getattr(module, tensor), name=tensor)
 
-    def initialize_weight(self, tensor):
-        pass
-
-    def initialize_bias(self, tensor):
-        pass
+    def initialize_tensor(self, tensor, name):
+        raise NotImplementedError

deeplay/initializers/kaiming.py

@@ -24,13 +24,20 @@ def __init__(
         targets: Tuple[Type[nn.Module], ...] = _kaiming_default_targets,
         mode: str = "fan_out",
         nonlinearity: str = "relu",
+        fill_bias: bool = True,
+        bias: float = 0.0,
     ):
         super().__init__(targets)
         self.mode = mode
         self.nonlinearity = nonlinearity
+        self.fill_bias = fill_bias
+        self.bias = bias
 
-    def initialize_weight(self, tensor):
-        nn.init.kaiming_normal_(tensor, mode=self.mode, nonlinearity=self.nonlinearity)
+    def initialize_tensor(self, tensor, name):
 
-    def initialize_bias(self, tensor):
-        tensor.data.fill_(0.0)
+        if name == "bias" and self.fill_bias:
+            tensor.data.fill_(self.bias)
+        else:
+            nn.init.kaiming_normal_(
+                tensor, mode=self.mode, nonlinearity=self.nonlinearity
+            )

deeplay/initializers/normal.py

@@ -28,8 +28,5 @@ def __init__(
         self.mean = mean
         self.std = std
 
-    def initialize_bias(self, tensor):
-        tensor.data.fill_(self.mean)
-
-    def initialize_weight(self, tensor):
-        tensor.data.normal_(self.mean, self.std)
+    def initialize_tensor(self, tensor, name):
+        tensor.data.normal_(mean=self.mean, std=self.std)

deeplay/module.py

@@ -1023,16 +1023,20 @@ def log_tensor(self, name: str, tensor: torch.Tensor):
         """
         self.logs[name] = tensor
 
-    def initialize(self, initializer):
+    def initialize(
+        self, initializer, tensors: Union[str, Tuple[str, ...]] = ("weight", "bias")
+    ):
+        if isinstance(tensors, str):
+            tensors = (tensors,)
         for module in self.modules():
             if isinstance(module, DeeplayModule):
-                module._initialize_after_build(initializer)
+                module._initialize_after_build(initializer, tensors)
             else:
-                initializer.initialize(module)
+                initializer.initialize(module, tensors)
 
     @after_build
-    def _initialize_after_build(self, initializer):
-        initializer.initialize(self)
+    def _initialize_after_build(self, initializer, tensors: Tuple[str, ...]):
+        initializer.initialize(self, tensors)
 
     @after_build
     def _validate_after_build(self):