Repository: mlcommons/GaNDLF

Issue:
Imagenet weights for GANDLF vgg/resnet etc for 2D images
**Is your feature request related to a problem? Please describe.**
It would be great to have pretrained imagenet weights in GANDLF for 2D images

**Describe the solution you'd like**
in the .yaml file, it would be great to have imagenet option set to true to work with pretrained weights if need rises.

**Describe alternatives you've considered**

**Additional context**
https://pytorch.org/vision/0.8/_modules/torchvision/models/vgg.html


Discussion:

Resolving patch:
diff --git a/GANDLF/parameterParsing.py b/GANDLF/parameterParsing.py
--- a/GANDLF/parameterParsing.py
+++ b/GANDLF/parameterParsing.py
@@ -54,6 +54,7 @@ def get_model(
             residualConnections=True,
         )
         divisibilityCheck_baseFilter = False
+
     elif modelname == "unet":
         model = unet(
             num_dimensions,
@@ -63,6 +64,7 @@ def get_model(
             final_convolution_layer=final_convolution_layer,
         )
         divisibilityCheck_baseFilter = False
+
     elif modelname == "fcn":
         model = fcn(
             num_dimensions,
@@ -74,6 +76,7 @@ def get_model(
         # not enough information to perform checking for this, yet
         divisibilityCheck_patch = False
         divisibilityCheck_baseFilter = False
+
     elif modelname == "uinc":
         model = uinc(
             num_dimensions,
@@ -82,6 +85,7 @@ def get_model(
             base_filters,
             final_convolution_layer=final_convolution_layer,
         )
+
     elif modelname == "msdnet":
         model = MSDNet(
             num_dimensions,
@@ -90,6 +94,67 @@ def get_model(
             base_filters,
             final_convolution_layer=final_convolution_layer,
         )
+
+    elif 'imagenet' in modelname: # these are generic imagenet-trained models and should be customized 
+        
+        if num_dimensions != 2:
+            sys.exit("ImageNet-trained models only work on 2D data")
+
+        divisibilityCheck_patch = False
+        divisibilityCheck_baseFilter = False
+
+        if 'batch_norm' in kwargs:
+            batch_norm = kwargs.get("batch_norm")
+        else:
+            batch_norm = True
+
+        if 'vgg11' in modelname:
+            if batch_norm:
+                model = torchvision.models.vgg11_bn(pretrained=True)
+            else:
+                model = torchvision.models.vgg11(pretrained=True)
+        elif 'vgg13' in modelname:
+            if batch_norm:
+                model = torchvision.models.vgg13_bn(pretrained=True)
+            else:
+                model = torchvision.models.vgg13(pretrained=True)
+        elif 'vgg16' in modelname:
+            if batch_norm:
+                model = torchvision.models.vgg16_bn(pretrained=True)
+            else:
+                model = torchvision.models.vgg16(pretrained=True)
+        elif 'vgg19' in modelname:
+            if batch_norm:
+                model = torchvision.models.vgg19_bn(pretrained=True)
+            else:
+                model = torchvision.models.vgg19(pretrained=True)
+        elif 'squeezenet1_0' in modelname:
+            model = torchvision.models.squeezenet1_0(pretrained=True)
+        elif 'squeezenet1_1' in modelname:
+            model = torchvision.models.squeezenet1_1(pretrained=True)
+        elif 'inceptionv3' in modelname:
+            model = torchvision.models.inception_v3(pretrained=True)
+        elif 'densenet121' in modelname:
+            model = torchvision.models.densenet121(pretrained=True)
+        elif 'densenet161' in modelname:
+            model = torchvision.models.densenet161(pretrained=True)
+        elif 'densenet169' in modelname:
+            model = torchvision.models.densenet169(pretrained=True)
+        elif 'densenet201' in modelname:
+            model = torchvision.models.densenet201(pretrained=True)
+        elif 'densenet264' in modelname:
+            model = torchvision.models.densenet264(pretrained=True)
+        elif 'resnet18' in modelname:
+            model = torchvision.models.resnet18(pretrained=True)
+        elif 'resnet34' in modelname:
+            model = torchvision.models.resnet34(pretrained=True)
+        elif 'resnet50' in modelname:
+            model = torchvision.models.resnet50(pretrained=True)
+        elif 'resnet101' in modelname:
+            model = torchvision.models.resnet101(pretrained=True)
+        elif 'resnet152' in modelname:
+            model = torchvision.models.resnet152(pretrained=True)
+        
     elif 'densenet' in modelname:
         if modelname == 'densenet121': # regressor/classifier network
             model = densenet.generate_model(model_depth=121,
@@ -121,12 +186,6 @@ def get_model(
         
         amp = False # this is not yet implemented for densenet 
 
-    # elif modelname == 'vgg16':
-    #     vgg_config = cfg['D']
-    #     num_final_features = vgg_config[-2]
-    #     divisibility_factor = Counter(vgg_config)["M"]
-    #     if patch_size[-1] == 1:
-    #         patch_size_altered = np.array(patch_size[:-1])
     elif 'vgg' in modelname: # common parsing for vgg
         if modelname == 'vgg11':
             vgg_config = cfg['A']
@@ -167,6 +226,7 @@ def get_model(
             num_classes,
             final_convolution_layer=final_convolution_layer,
         )
+
     elif modelname == "brain_age":
         if num_dimensions != 2:
             sys.exit("Brain Age predictions only works on 2D data")


Test patch:
diff --git a/testing/config_classification.yaml b/testing/config_classification.yaml
--- a/testing/config_classification.yaml
+++ b/testing/config_classification.yaml
@@ -6,6 +6,8 @@ data_preprocessing:
 in_memory: false
 learning_rate: 0.001
 loss_function: mse
+metrics:
+  - mse
 model:
   amp: false
   architecture: vgg16
diff --git a/testing/config_regression.yaml b/testing/config_regression.yaml
--- a/testing/config_regression.yaml
+++ b/testing/config_regression.yaml
@@ -6,6 +6,8 @@ data_preprocessing:
 in_memory: false
 learning_rate: 0.001
 loss_function: mse
+metrics:
+  - mse
 model:
   amp: false
   architecture: densenet121
diff --git a/testing/config_segmentation.yaml b/testing/config_segmentation.yaml
--- a/testing/config_segmentation.yaml
+++ b/testing/config_segmentation.yaml
@@ -15,6 +15,8 @@ model:
     amp: False, # Set if you want to use Automatic Mixed Precision for your operations or not - options: True, False
     # n_channels: 3, # set the input channels - useful when reading RGB or images that have vectored pixel types
   }
+metrics:
+  - dice
 # Patch size during training - 2D patch for breast images since third dimension is not patched 
 patch_size: [128,128]
 # Number of epochs
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -110,7 +110,6 @@ def test_train_segmentation_rad_2d(device):
   parameters['psize'] = patch_size['2D']
   parameters['model']['dimension'] = 2
   parameters['model']['class_list'] = [0,255]
-  parameters['metrics'] = ['dice']
   parameters['model']['amp'] = True
   parameters['model']['num_channels'] = len(headers["channelHeaders"])
   # read and initialize parameters for specific data dimension
@@ -132,7 +131,6 @@ def test_train_segmentation_rad_3d(device):
   parameters['psize'] = patch_size['3D']
   parameters['model']['dimension'] = 3
   parameters['model']['class_list'] = [0,1]
-  parameters['metrics'] = ['dice']
   parameters['model']['amp'] = True
   parameters['model']['num_channels'] = len(headers["channelHeaders"])
   # loop through selected models and train for single epoch
@@ -150,7 +148,6 @@ def test_train_regression_rad_2d(device):
   parameters['patch_size'] = patch_size['2D']
   parameters['psize'] = patch_size['2D']
   parameters['model']['dimension'] = 2
-  parameters['metrics'] = ['mse']
   parameters['model']['amp'] = True
   # read and parse csv 
   training_data, headers = parseTrainingCSV(inputDir + '/train_2d_rad_regression.csv')
@@ -172,7 +169,6 @@ def test_train_regression_rad_3d(device):
   parameters['patch_size'] = patch_size['3D']
   parameters['psize'] = patch_size['3D']
   parameters['model']['dimension'] = 3
-  parameters['metrics'] = ['mse']
   parameters['model']['amp'] = True
   # read and parse csv 
   training_data, headers = parseTrainingCSV(inputDir + '/train_3d_rad_regression.csv')
@@ -194,7 +190,6 @@ def test_train_classification_rad_2d(device):
   parameters['patch_size'] = patch_size['2D']
   parameters['psize'] = patch_size['2D']
   parameters['model']['dimension'] = 2
-  parameters['metrics'] = ['mse']
   parameters['model']['amp'] = True
   # read and parse csv 
   training_data, headers = parseTrainingCSV(inputDir + '/train_2d_rad_classification.csv')
@@ -215,7 +210,6 @@ def test_train_classification_rad_3d(device):
   parameters['patch_size'] = patch_size['3D']
   parameters['psize'] = patch_size['3D']
   parameters['model']['dimension'] = 3
-  parameters['metrics'] = ['mse']
   parameters['model']['amp'] = True
   # read and parse csv 
   training_data, headers = parseTrainingCSV(inputDir + '/train_3d_rad_classification.csv')