tensorboard notebook fix & loading of pre-trained models fix (#397)

Co-authored-by: Max Kaznady <max.kaznady@gmail.com>

cv_lib/cv_lib/segmentation/models/patch_deconvnet.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-
+import torch
 import torch.nn as nn
 
 
@@ -304,5 +304,9 @@ def get_seg_model(cfg, **kwargs):
         cfg.MODEL.IN_CHANNELS == 1
     ), f"Patch deconvnet is not implemented to accept {cfg.MODEL.IN_CHANNELS} channels. Please only pass 1 for cfg.MODEL.IN_CHANNELS"
     model = patch_deconvnet(n_classes=cfg.DATASET.NUM_CLASSES)
-
+    # load the pre-trained model
+    if "PRETRAINED" in cfg.MODEL.keys():
+        trained_model = torch.load(cfg.MODEL.PRETRAINED)
+        trained_model = {k.replace("module.", ""): v for (k, v) in trained_model.items()}
+        model.load_state_dict(trained_model, strict=True)
     return model

cv_lib/cv_lib/segmentation/models/patch_deconvnet_skip.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-
+import torch
 import torch.nn as nn
 
 
@@ -304,5 +304,9 @@ def get_seg_model(cfg, **kwargs):
         cfg.MODEL.IN_CHANNELS == 1
     ), f"Patch deconvnet is not implemented to accept {cfg.MODEL.IN_CHANNELS} channels. Please only pass 1 for cfg.MODEL.IN_CHANNELS"
     model = patch_deconvnet_skip(n_classes=cfg.DATASET.NUM_CLASSES)
-
+    # load the pre-trained model
+    if "PRETRAINED" in cfg.MODEL.keys():
+        trained_model = torch.load(cfg.MODEL.PRETRAINED)
+        trained_model = {k.replace("module.", ""): v for (k, v) in trained_model.items()}
+        model.load_state_dict(trained_model, strict=True)
     return model

cv_lib/cv_lib/segmentation/models/resnet_unet.py

@@ -367,4 +367,9 @@ def get_seg_model(cfg, **kwargs):
         cfg.MODEL.IN_CHANNELS == 3
     ), f"SEResnet Unet deconvnet is not implemented to accept {cfg.MODEL.IN_CHANNELS} channels. Please only pass 3 for cfg.MODEL.IN_CHANNELS"
     model = Res34Unetv4(n_classes=cfg.DATASET.NUM_CLASSES)
+    # load the pre-trained model
+    if "PRETRAINED" in cfg.MODEL.keys():
+        trained_model = torch.load(cfg.MODEL.PRETRAINED)
+        trained_model = {k.replace("module.", ""): v for (k, v) in trained_model.items()}
+        model.load_state_dict(trained_model, strict=True)
     return model

cv_lib/cv_lib/segmentation/models/section_deconvnet.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-
+import torch
 import torch.nn as nn
 
 
@@ -304,5 +304,9 @@ def get_seg_model(cfg, **kwargs):
         cfg.MODEL.IN_CHANNELS == 1
     ), f"Section deconvnet is not implemented to accept {cfg.MODEL.IN_CHANNELS} channels. Please only pass 1 for cfg.MODEL.IN_CHANNELS"
     model = section_deconvnet(n_classes=cfg.DATASET.NUM_CLASSES)
-
+    # load the pre-trained model
+    if "PRETRAINED" in cfg.MODEL.keys():
+        trained_model = torch.load(cfg.MODEL.PRETRAINED)
+        trained_model = {k.replace("module.", ""): v for (k, v) in trained_model.items()}
+        model.load_state_dict(trained_model, strict=True)
     return model

cv_lib/cv_lib/segmentation/models/section_deconvnet_skip.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-
+import torch
 import torch.nn as nn
 
 
@@ -304,5 +304,9 @@ def get_seg_model(cfg, **kwargs):
         cfg.MODEL.IN_CHANNELS == 1
     ), f"Section deconvnet is not implemented to accept {cfg.MODEL.IN_CHANNELS} channels. Please only pass 1 for cfg.MODEL.IN_CHANNELS"
     model = section_deconvnet_skip(n_classes=cfg.DATASET.NUM_CLASSES)
-
+    # load the pre-trained model
+    if "PRETRAINED" in cfg.MODEL.keys():
+        trained_model = torch.load(cfg.MODEL.PRETRAINED)
+        trained_model = {k.replace("module.", ""): v for (k, v) in trained_model.items()}
+        model.load_state_dict(trained_model, strict=True)
     return model

cv_lib/cv_lib/segmentation/models/unet.py

@@ -113,5 +113,9 @@ def forward(self, x):
 
 def get_seg_model(cfg, **kwargs):
     model = UNet(cfg.MODEL.IN_CHANNELS, cfg.DATASET.NUM_CLASSES)
-
+    # load the pre-trained model
+    if "PRETRAINED" in cfg.MODEL.keys():
+        trained_model = torch.load(cfg.MODEL.PRETRAINED)
+        trained_model = {k.replace("module.", ""): v for (k, v) in trained_model.items()}
+        model.load_state_dict(trained_model, strict=True)
     return model

environment/anaconda/local/environment.yml

@@ -12,7 +12,7 @@ dependencies:
   - torchvision>=0.5.0
   - pandas==0.25.3
   - scikit-learn==0.21.3
-  - tensorflow==2.1.0
+  - tensorflow==2.1
   - opt-einsum>=2.3.2
   - tqdm==4.39.0
   - itkwidgets==0.23.1

interpretation/deepseismic_interpretation/models/texture_net.py

@@ -156,4 +156,9 @@ def getFeatures(self, x, layer_no):
 
 def get_seg_model(cfg, **kwargs):
     model = TextureNet(n_classes=cfg.DATASET.NUM_CLASSES)
+    # load the pre-trained model
+    if "PRETRAINED" in cfg.MODEL.keys():
+        trained_model = torch.load(cfg.MODEL.PRETRAINED)
+        trained_model = {k.replace("module.", ""): v for (k, v) in trained_model.items()}
+        model.load_state_dict(trained_model, strict=True)
     return model