Repository: mlcommons/GaNDLF

Issue:
Memory issue for increasing pytest test cases
**Describe your question/problem**
I tried to incorporate TransUNet, but the test cases are failing with exit code 137. When I comment out all test cases except for the 2 relevant to TransUNet, it completes. I have seen some instances on google where pytest memory accumulates as the test suite grows. The code is in a draft PR.


Discussion:

Resolving patch:
diff --git a/GANDLF/models/__init__.py b/GANDLF/models/__init__.py
--- a/GANDLF/models/__init__.py
+++ b/GANDLF/models/__init__.py
@@ -32,6 +32,7 @@
 from .MSDNet import MSDNet
 from .brain_age import brainage
 from .unetr import unetr
+from .transunet import transunet
 
 # defining dict for models - key is the string and the value is the transform object
 global_models_dict = {
@@ -54,6 +55,7 @@
     "light_resunet": light_resunet,
     "light_resunet_multilayer": light_resunet_multilayer,
     "unetr": unetr,
+    "transunet": transunet,
     "fcn": fcn,
     "uinc": uinc,
     "vgg": vgg19,
diff --git a/GANDLF/models/transunet.py b/GANDLF/models/transunet.py
new file mode 100644
--- /dev/null
+++ b/GANDLF/models/transunet.py
@@ -0,0 +1,223 @@
+# -*- coding: utf-8 -*-
+"""
+Implementation of TransUNet
+"""
+from torch.nn import ModuleList
+
+from GANDLF.models.seg_modules.DownsamplingModule import DownsamplingModule
+from GANDLF.models.seg_modules.EncodingModule import EncodingModule
+from GANDLF.models.seg_modules.DecodingModule import DecodingModule
+from GANDLF.models.seg_modules.UpsamplingModule import UpsamplingModule
+from GANDLF.models.seg_modules.in_conv import in_conv
+from GANDLF.models.seg_modules.out_conv import out_conv
+from .modelBase import ModelBase
+import sys
+from GANDLF.utils.generic import checkPatchDimensions
+from GANDLF.models.seg_modules.Interpolate import Interpolate
+import torch
+import torch.nn as nn
+from torch.nn import ModuleList
+import numpy as np
+import math
+from .unetr import _Transformer
+
+
+class _DecoderCUP(nn.Sequential):
+    def __init__(self, in_feats, out_feats, Norm, Conv, Upsample):
+        super().__init__()
+
+        self.conv = Conv(
+            in_channels=in_feats,
+            out_channels=out_feats,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=False,
+        )
+
+        self.norm = Norm(out_feats)
+        self.relu = nn.ReLU(inplace=True)
+        self.upsample = Upsample
+
+    def forward(self, x1, x2):
+        if x2 is not None:
+            x1 = torch.cat([x1, x2], dim=1)
+        x = self.conv(x1)
+        x = self.norm(x)
+        x = self.relu(x)
+        x = self.upsample(x)
+
+        return x
+
+
+class transunet(ModelBase):
+    """
+    This is the TransUNet architecture : https://doi.org/10.48550/arXiv.2102.04306. The Downsampling, Encoding modules
+    are defined in the seg_modules file. These smaller modules are basically defined by 2 parameters, the input channels (filters) and the output channels (filters),
+    and some other hyperparameters, which remain constant all the modules. For more details on the smaller modules please have a look at the seg_modules file.
+    """
+
+    def __init__(
+        self,
+        parameters: dict,
+    ):
+        super(transunet, self).__init__(parameters)
+
+        if not ("depth" in parameters["model"]):
+            parameters["model"]["depth"] = 4
+            print("Default depth set to 4.")
+
+        patch_check = checkPatchDimensions(
+            parameters["patch_size"], numlay=parameters["model"]["depth"]
+        )
+
+        if patch_check != parameters["model"]["depth"] and patch_check >= 2:
+            print(
+                "The patch size is not large enough for desired depth. It is expected that each dimension of the patch size is divisible by 2^i, where i is in a integer greater than or equal to 2. Only the first %d layers will run."
+                % patch_check
+            )
+        elif patch_check < 2:
+            sys.exit(
+                "The patch size is not large enough for desired depth. It is expected that each dimension of the patch size is divisible by 2^i, where i is in a integer greater than or equal to 2."
+            )
+
+        if not ("num_heads" in parameters["model"]):
+            parameters["model"]["num_heads"] = 12
+            print(
+                "Default number of heads in multi-head self-attention (MSA) set to 12."
+            )
+
+        if not ("embed_dim" in parameters["model"]):
+            parameters["model"]["embed_dim"] = 768
+            print("Default size of embedded dimension set to 768.")
+
+        self.depth = patch_check
+
+        if self.n_dimensions == 2:
+            self.img_size = parameters["patch_size"][0:2]
+            self.upsample = nn.Upsample(
+                scale_factor=2, mode="bilinear", align_corners=True
+            )
+        elif self.n_dimensions == 3:
+            self.img_size = parameters["patch_size"]
+            self.upsample = nn.Upsample(
+                scale_factor=2, mode="trilinear", align_corners=True
+            )
+
+        self.num_layers = 3 * self.depth  # number of transformer layers
+        self.out_layers = [self.num_layers - 1]
+
+        self.num_heads = parameters["model"]["num_heads"]
+        self.embed_size = parameters["model"]["embed_dim"]
+
+        if self.embed_size % self.num_heads != 0:
+            sys.exit(
+                "The embedding dimension must be divisible by the number of self-attention heads."
+            )
+
+        self.patch_dim = [i // 2 ** (self.depth) for i in self.img_size]
+
+        self.ins = in_conv(
+            input_channels=self.n_channels,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            dropout=self.Dropout,
+            norm=self.Norm,
+        )
+
+        self.ds = ModuleList([])
+        self.en = ModuleList([])
+        self.de = ModuleList([])
+
+        for i_lay in range(0, self.depth):
+            self.ds.append(
+                DownsamplingModule(
+                    input_channels=self.base_filters * 2 ** (i_lay),
+                    output_channels=self.base_filters * 2 ** (i_lay + 1),
+                    conv=self.Conv,
+                    norm=self.Norm,
+                )
+            )
+
+            self.de.append(
+                _DecoderCUP(
+                    in_feats=2 * self.base_filters * 2 ** (i_lay + 1),
+                    out_feats=self.base_filters * 2 ** (i_lay),
+                    Conv=self.Conv,
+                    Norm=self.Norm,
+                    Upsample=self.upsample,
+                )
+            )
+
+            self.en.append(
+                EncodingModule(
+                    input_channels=self.base_filters * 2 ** (i_lay + 1),
+                    output_channels=self.base_filters * 2 ** (i_lay + 1),
+                    conv=self.Conv,
+                    dropout=self.Dropout,
+                    norm=self.Norm,
+                )
+            )
+
+        self.transformer = _Transformer(
+            img_size=[i // 2 ** (self.depth) for i in self.img_size],
+            patch_size=1,
+            in_feats=self.base_filters * 2**self.depth,
+            embed_size=self.embed_size,
+            num_heads=self.num_heads,
+            mlp_dim=2048,
+            num_layers=self.num_layers,
+            out_layers=self.out_layers,
+            Conv=self.Conv,
+            Norm=self.Norm,
+        )
+
+        self.transCUP = _DecoderCUP(
+            in_feats=self.embed_size,
+            out_feats=self.base_filters * 2 ** (self.depth - 1),
+            Conv=self.Conv,
+            Norm=self.Norm,
+            Upsample=self.upsample,
+        )
+
+        # TODO: conv 3x3 --> ReLU --> outconv
+        self.out = out_conv(
+            input_channels=self.base_filters,
+            output_channels=self.n_classes,
+            conv=self.Conv,
+            norm=self.Norm,
+            final_convolution_layer=self.final_convolution_layer,
+            sigmoid_input_multiplier=self.sigmoid_input_multiplier,
+        )
+
+    def forward(self, x):
+        """
+        Parameters
+        ----------
+        x : Tensor
+            Should be a 5D Tensor as [batch_size, channels, x_dims, y_dims, z_dims].
+
+        Returns
+        -------
+        x : Tensor
+            Returns a 5D Output Tensor as [batch_size, n_classes, x_dims, y_dims, z_dims].
+
+        """
+        y = []
+        y.append(self.ins(x))
+
+        # [downsample --> encode] x num layers
+        for i in range(0, self.depth):
+            temp = self.ds[i](y[i])
+            y.append(self.en[i](temp))
+
+        x = self.transformer(y[-1])[-1]
+        x = x.transpose(-1, -2).view(-1, self.embed_size, *self.patch_dim)
+        x = self.transCUP(x, None)
+
+        # [upsample --> encode] x num layers
+        for i in range(self.depth - 1, 0, -1):
+            x = self.de[i - 1](x, y[i])
+
+        x = self.out(x)
+        return x
diff --git a/GANDLF/models/unetr.py b/GANDLF/models/unetr.py
--- a/GANDLF/models/unetr.py
+++ b/GANDLF/models/unetr.py
@@ -316,7 +316,9 @@ def __init__(
                 sys.exit("The inner patch size must be a power of 2.")
 
         self.patch_size = parameters["model"]["inner_patch_size"]
+
         self.depth = int(np.log2(self.patch_size))
+
         patch_check = checkPatchDimensions(parameters["patch_size"], self.depth)
 
         if patch_check != self.depth and patch_check >= 2:
@@ -346,8 +348,15 @@ def __init__(
 
         self.num_layers = 3 * self.depth  # number of transformer layers
         self.out_layers = np.arange(2, self.num_layers, 3)
+
         self.num_heads = parameters["model"]["num_heads"]
         self.embed_size = parameters["model"]["embed_dim"]
+
+        if self.embed_size % self.num_heads != 0:
+            sys.exit(
+                "The embedding dimension must be divisible by the number of self-attention heads."
+            )
+
         self.patch_dim = [i // self.patch_size for i in self.img_size]
 
         if not all([i % self.patch_size == 0 for i in self.img_size]):


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -81,7 +81,7 @@
 """
 
 
-def test_download_data():
+def test_generic_download_data():
     print("00: Downloading the sample data")
     urlToDownload = (
         "https://upenn.box.com/shared/static/y8162xkq1zz5555ye3pwadry2m2e39bs.zip"
@@ -104,7 +104,7 @@ def test_download_data():
     print("passed")
 
 
-def test_constructTrainingCSV():
+def test_generic_constructTrainingCSV():
     print("01: Constructing training CSVs")
     # inputDir = os.path.normpath('./testing/data')
     # delete previous csv files
@@ -567,7 +567,7 @@ def test_train_resume_inference_classification_rad_3d(device):
     print("passed")
 
 
-def test_inference_optimize_classification_rad_3d(device):
+def test_train_inference_optimize_classification_rad_3d(device):
     print("13: Starting 3D Rad segmentation tests for optimization")
     # read and initialize parameters for specific data dimension
     parameters = parseConfig(
@@ -608,7 +608,7 @@ def test_inference_optimize_classification_rad_3d(device):
     print("passed")
 
 
-def test_inference_optimize_segmentation_rad_2d(device):
+def test_train_inference_optimize_segmentation_rad_2d(device):
     print("14: Starting 2D Rad segmentation tests for optimization")
     # read and parse csv
     parameters = parseConfig(
@@ -652,7 +652,7 @@ def test_inference_optimize_segmentation_rad_2d(device):
     print("passed")
 
 
-def test_inference_classification_with_logits_single_fold_rad_3d(device):
+def test_train_inference_classification_with_logits_single_fold_rad_3d(device):
     print("15: Starting 3D Rad classification tests for single fold logits inference")
     # read and initialize parameters for specific data dimension
     parameters = parseConfig(
@@ -714,7 +714,7 @@ def test_inference_classification_with_logits_single_fold_rad_3d(device):
     print("passed")
 
 
-def test_inference_classification_with_logits_multiple_folds_rad_3d(device):
+def test_train_inference_classification_with_logits_multiple_folds_rad_3d(device):
     print("16: Starting 3D Rad classification tests for multi-fold logits inference")
     # read and initialize parameters for specific data dimension
     parameters = parseConfig(
@@ -756,7 +756,7 @@ def test_inference_classification_with_logits_multiple_folds_rad_3d(device):
     print("passed")
 
 
-def test_scheduler_classification_rad_2d(device):
+def test_train_scheduler_classification_rad_2d(device):
     print("17: Starting 2D Rad segmentation tests for scheduler")
     # read and initialize parameters for specific data dimension
     # loop through selected models and train for single epoch
@@ -803,7 +803,7 @@ def test_scheduler_classification_rad_2d(device):
     print("passed")
 
 
-def test_optimizer_classification_rad_2d(device):
+def test_train_optimizer_classification_rad_2d(device):
     print("18: Starting 2D Rad classification tests for optimizer")
     # read and initialize parameters for specific data dimension
     parameters = parseConfig(
@@ -877,7 +877,7 @@ def test_clip_train_classification_rad_3d(device):
     print("passed")
 
 
-def test_normtype_train_segmentation_rad_3d(device):
+def test_train_normtype_segmentation_rad_3d(device):
     print("20: Starting 3D Rad segmentation tests for normtype")
     # read and initialize parameters for specific data dimension
     # read and parse csv
@@ -920,7 +920,7 @@ def test_normtype_train_segmentation_rad_3d(device):
         print("passed")
 
 
-def test_metrics_segmentation_rad_2d(device):
+def test_train_metrics_segmentation_rad_2d(device):
     print("21: Starting 2D Rad segmentation tests for metrics")
     # read and parse csv
     parameters = parseConfig(
@@ -953,7 +953,7 @@ def test_metrics_segmentation_rad_2d(device):
     print("passed")
 
 
-def test_metrics_regression_rad_2d(device):
+def test_train_metrics_regression_rad_2d(device):
     print("22: Starting 2D Rad regression tests for metrics")
     # read and parse csv
     parameters = parseConfig(
@@ -985,7 +985,7 @@ def test_metrics_regression_rad_2d(device):
     print("passed")
 
 
-def test_losses_segmentation_rad_2d(device):
+def test_train_losses_segmentation_rad_2d(device):
     print("23: Starting 2D Rad segmentation tests for losses")
     # read and parse csv
     parameters = parseConfig(
@@ -1023,7 +1023,7 @@ def test_losses_segmentation_rad_2d(device):
     print("passed")
 
 
-def test_config_read():
+def test_generic_config_read():
     print("24: Starting testing reading configuration")
     # read and parse csv
     file_config_temp = os.path.join(outputDir, "config_segmentation_temp.yaml")
@@ -1117,7 +1117,7 @@ def test_config_read():
     print("passed")
 
 
-def test_cli_function_preprocess():
+def test_generic_cli_function_preprocess():
     print("25: Starting testing cli function preprocess")
     file_config = os.path.join(testingDir, "config_segmentation.yaml")
     sanitize_outputDir()
@@ -1155,7 +1155,7 @@ def test_cli_function_preprocess():
     print("passed")
 
 
-def test_cli_function_mainrun(device):
+def test_generic_cli_function_mainrun(device):
     print("26: Starting testing cli function main_run")
     parameters = parseConfig(
         testingDir + "/config_segmentation.yaml", version_check_flag=False
@@ -1246,7 +1246,7 @@ def test_dataloader_construction_train_segmentation_3d(device):
     print("passed")
 
 
-def test_preprocess_functions():
+def test_generic_preprocess_functions():
     print("28: Starting testing preprocessing functions")
     # initialize an input which has values between [-1,1]
     # checking tensor with last dimension of size 1
@@ -1372,7 +1372,7 @@ def test_preprocess_functions():
     print("passed")
 
 
-def test_augmentation_functions():
+def test_generic_augmentation_functions():
     print("29: Starting testing augmentation functions")
     params_all_preprocessing_and_augs = parseConfig(
         os.path.join(baseConfigDir, "config_all_options.yaml")
@@ -1414,7 +1414,7 @@ def test_augmentation_functions():
     print("passed")
 
 
-def test_checkpointing_segmentation_rad_2d(device):
+def test_train_checkpointing_segmentation_rad_2d(device):
     print("30: Starting 2D Rad segmentation tests for metrics")
     # read and parse csv
     parameters = parseConfig(
@@ -1466,7 +1466,7 @@ def test_checkpointing_segmentation_rad_2d(device):
     print("passed")
 
 
-def test_model_patch_divisibility():
+def test_generic_model_patch_divisibility():
     print("31: Starting patch divisibility tests")
     parameters = parseConfig(
         testingDir + "/config_segmentation.yaml", version_check_flag=False
@@ -1499,7 +1499,7 @@ def test_model_patch_divisibility():
     print("passed")
 
 
-def test_one_hot_logic():
+def test_generic_one_hot_logic():
     print("32: Starting one hot logic tests")
     random_array = np.random.randint(5, size=(20, 20, 20))
     img = sitk.GetImageFromArray(random_array)
@@ -1573,7 +1573,7 @@ def test_one_hot_logic():
     print("passed")
 
 
-def test_anonymizer():
+def test_generic_anonymizer():
     print("33: Starting anomymizer tests")
     input_file = get_testdata_file("MR_small.dcm")
 
@@ -1777,7 +1777,7 @@ def test_train_inference_classification_histology_2d(device):
     print("passed")
 
 
-def test_unet_layerchange_2d(device):
+def test_train_unet_layerchange_rad_2d(device):
     # test case to up code coverage --> test decreasing allowed layers for unet
     print("36: Starting 2D Rad segmentation tests for normtype")
     # read and parse csv
@@ -1824,7 +1824,7 @@ def test_unet_layerchange_2d(device):
     print("passed")
 
 
-def test_train_segmentation_unetr_3d(device):
+def test_train_segmentation_unetr_rad_3d(device):
     print("37: Testing UNETR for 3D segmentation")
     parameters = parseConfig(
         testingDir + "/config_segmentation.yaml", version_check_flag=False
@@ -1879,7 +1879,7 @@ def test_train_segmentation_unetr_3d(device):
     print("passed")
 
 
-def test_train_segmentation_unetr_2d(device):
+def test_train_segmentation_unetr_rad_2d(device):
     print("38: Testing UNETR for 2D segmentation")
     parameters = parseConfig(
         testingDir + "/config_segmentation.yaml", version_check_flag=False
@@ -1914,3 +1914,110 @@ def test_train_segmentation_unetr_2d(device):
         )
 
     print("passed")
+
+
+def test_train_segmentation_transunet_rad_2d(device):
+    print("39: Testing TransUNet for 2D segmentation")
+    parameters = parseConfig(
+        testingDir + "/config_segmentation.yaml", version_check_flag=False
+    )
+    training_data, parameters["headers"] = parseTrainingCSV(
+        inputDir + "/train_2d_rad_segmentation.csv"
+    )
+    parameters["model"]["architecture"] = "transunet"
+    parameters["patch_size"] = [128, 128, 1]
+    parameters["model"]["dimension"] = 2
+
+    with pytest.raises(BaseException) as e_info:
+        parameters["model"]["num_heads"] = 6
+        parameters["model"]["embed_dim"] = 64
+        global_models_dict[parameters["model"]["architecture"]](parameters=parameters)
+
+    with pytest.raises(BaseException) as e_info:
+        parameters["model"]["num_heads"] = 3
+        parameters["model"]["embed_dim"] = 50
+        global_models_dict[parameters["model"]["architecture"]](parameters=parameters)
+
+    parameters["model"]["embed_dim"] = 64
+    parameters["model"]["depth"] = 2
+    parameters["model"]["class_list"] = [0, 255]
+    parameters["model"]["num_heads"] = 8
+    parameters["model"]["amp"] = True
+    parameters["model"]["num_channels"] = 3
+    parameters = populate_header_in_parameters(parameters, parameters["headers"])
+    # loop through selected models and train for single epoch
+    parameters["model"]["norm_type"] = "batch"
+    parameters["nested_training"]["testing"] = -5
+    parameters["nested_training"]["validation"] = -5
+    if os.path.isdir(outputDir):
+        shutil.rmtree(outputDir)  # overwrite previous results
+    sanitize_outputDir()
+    TrainingManager(
+        dataframe=training_data,
+        outputDir=outputDir,
+        parameters=parameters,
+        device=device,
+        resume=False,
+        reset=True,
+    )
+
+    print("passed")
+
+
+def test_train_segmentation_transunet_rad_3d(device):
+    print("40: Testing TransUNet for 3D segmentation")
+    parameters = parseConfig(
+        testingDir + "/config_segmentation.yaml", version_check_flag=False
+    )
+    training_data, parameters["headers"] = parseTrainingCSV(
+        inputDir + "/train_3d_rad_segmentation.csv"
+    )
+    parameters["model"]["architecture"] = "transunet"
+    parameters["patch_size"] = [4, 4, 4]
+    parameters["model"]["dimension"] = 3
+
+    # this assertion should fail
+    with pytest.raises(BaseException) as e_info:
+        global_models_dict[parameters["model"]["architecture"]](parameters=parameters)
+
+    parameters["model"]["dimension"] = 3
+    parameters["patch_size"] = [32, 32, 32]
+
+    with pytest.raises(BaseException) as e_info:
+        parameters["model"]["depth"] = 1
+        global_models_dict[parameters["model"]["architecture"]](parameters=parameters)
+
+    with pytest.raises(BaseException) as e_info:
+        parameters["model"]["num_heads"] = 6
+        parameters["model"]["embed_dim"] = 64
+        global_models_dict[parameters["model"]["architecture"]](parameters=parameters)
+
+    with pytest.raises(BaseException) as e_info:
+        parameters["model"]["num_heads"] = 3
+        parameters["model"]["embed_dim"] = 50
+        global_models_dict[parameters["model"]["architecture"]](parameters=parameters)
+
+    parameters["model"]["num_heads"] = 8
+    parameters["model"]["embed_dim"] = 64
+    parameters["model"]["depth"] = 2
+    parameters["model"]["class_list"] = [0, 255]
+    parameters["model"]["amp"] = True
+    parameters["model"]["num_channels"] = len(parameters["headers"]["channelHeaders"])
+    parameters = populate_header_in_parameters(parameters, parameters["headers"])
+    # loop through selected models and train for single epoch
+    parameters["model"]["norm_type"] = "batch"
+    parameters["nested_training"]["testing"] = -5
+    parameters["nested_training"]["validation"] = -5
+    if os.path.isdir(outputDir):
+        shutil.rmtree(outputDir)  # overwrite previous results
+    sanitize_outputDir()
+    TrainingManager(
+        dataframe=training_data,
+        outputDir=outputDir,
+        parameters=parameters,
+        device=device,
+        resume=False,
+        reset=True,
+    )
+
+    print("passed")