Offer config parameter to initialize model weights

funkelab · Oct 12, 2023 · 4b17700 · 4b17700
1 parent 255fa32
commit 4b17700
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/cellulus/train.py b/cellulus/train.py
@@ -58,9 +58,10 @@ def train(experiment_config):
     model = model.to(device)
 
     # initialize model weights
-    # for _name, layer in model.named_modules():
-    #    if isinstance(layer, torch.nn.modules.conv._ConvNd):
-    #        torch.nn.init.kaiming_normal_(layer.weight, nonlinearity="relu")
+    if model_config.initialize:
+        for _name, layer in model.named_modules():
+            if isinstance(layer, torch.nn.modules.conv._ConvNd):
+                torch.nn.init.kaiming_normal_(layer.weight, nonlinearity="relu")
 
     # set loss
     criterion = get_loss(