Added custom modules

Author: claesnl

rhtorch/models/modules.py

@@ -2,44 +2,9 @@
 import torch
 import torch.nn as nn
 import pytorch_lightning as pl
-import rhtorch
 import torchmetrics as tm
 import math
-import pkgutil
-import os
-import sys
-import importlib
-
-
-def recursive_find_python_class(name, folder=None, current_module="rhtorch.models"):
-
-    # Set default search path to root modules
-    if folder is None:
-        folder = [os.path.join(rhtorch.__path__[0], 'models')]
-
-    tr = None
-    for importer, modname, ispkg in pkgutil.iter_modules(folder):
-        if not ispkg:
-            m = importlib.import_module(current_module + '.' + modname)
-            if hasattr(m, name):
-                tr = getattr(m, name)
-                break
-
-    if tr is None:
-        for importer, modname, ispkg in pkgutil.iter_modules(folder):
-            if ispkg:
-                next_current_module = current_module + '.' + modname
-                tr = recursive_find_python_class(name, folder=[os.path.join(
-                    folder[0], modname)], current_module=next_current_module)
-
-            if tr is not None:
-                break
-
-    if tr is None:
-        sys.exit(f"Could not find module {name}")
-
-    return tr
-
+from rhtorch.utilities.modules import recursive_find_python_class 
 
 class LightningAE(pl.LightningModule):
     def __init__(self, hparams, in_shape=(2, 128, 128, 128)):

rhtorch/torch_training.py

@@ -13,7 +13,7 @@
 from rhtorch.models import modules
 from rhtorch.callbacks import plotting
 from rhtorch.config_utils import UserConfig
-
+from rhtorch.utilities.modules import recursive_find_python_class
 
 def main():
     import argparse
@@ -65,8 +65,8 @@ def main():
 
     # define lightning module
     shape_in = data_train.data_shape_in
-    module = getattr(modules, configs['module'])
-    model = module(configs, shape_in)
+    module = recursive_find_python_class( configs['module'] )
+    model = module(configs, shape_in) # Should also be changed to custom arguments (**configs)
 
     # transfer learning setup
     if 'pretrained generator' in configs and configs['pretrained_generator']:
@@ -79,8 +79,6 @@ def main():
                     pretrained_model_path, hparams=configs, in_shape=shape_in, strict=False)
             elif pretrained_model_path.name.endswith(".pt"):
                 ckpt = torch.load(pretrained_model_path)
-                # OBS, the 'state_dict' is not set during save?
-                # What if we are to save multiple models used later for pretrain? (e.g. a GAN with 3 networks?)
                 pretrained_model = ckpt['state_dict'] if 'state_dict' in ckpt else ckpt
                 model.load_state_dict(pretrained_model, strict=False)
             else:
@@ -89,6 +87,7 @@ def main():
             raise FileNotFoundError(
                 "Model file not found. Check your path in config file.")
 
+        # This should be more generic. What if we are to only freeze some of the layers?
         if 'freeze_encoder' in configs and configs['freeze_encoder']:
             message += "with frozen encoder."
             model.encoder.freeze()

rhtorch/utilities/__init__.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue May 25 19:40:15 2021
+
+@author: clad0003
+"""
+
+import rhtorch
+import pkgutil
+import os
+import sys
+import importlib
+
+def recursive_find_python_class(name, folder=None, current_module="rhtorch.models"):
+
+    # Set default search path to root modules
+    if folder is None:
+        folder = [os.path.join(rhtorch.__path__[0], 'models')]
+
+    tr = None
+    for importer, modname, ispkg in pkgutil.iter_modules(folder):
+        if not ispkg:
+            m = importlib.import_module(current_module + '.' + modname)
+            if hasattr(m, name):
+                tr = getattr(m, name)
+                break
+
+    if tr is None:
+        for importer, modname, ispkg in pkgutil.iter_modules(folder):
+            if ispkg:
+                next_current_module = current_module + '.' + modname
+                tr = recursive_find_python_class(name, folder=[os.path.join(
+                    folder[0], modname)], current_module=next_current_module)
+
+            if tr is not None:
+                break
+
+    if tr is None:
+        sys.exit(f"Could not find module {name}")
+
+    return tr