Repository: mlcommons/GaNDLF

Issue:
`light_unet` is not working via unit tests
**Describe the bug**
When `light_unet` is added to the `all_modes_segmentation` [[ref](https://github.com/CBICA/GaNDLF/blob/master/testing/test_full.py#L12)], the tests fail.

**To Reproduce**
Steps to reproduce the behavior:
1. Add `light_unet` to `all_modes_segmentation` [[ref](https://github.com/CBICA/GaNDLF/blob/master/testing/test_full.py#L12)]
2. Run `pytest` [[ref](https://cbica.github.io/GaNDLF/extending#run-tests)]
3. See error

**Expected behavior**
It should pass.

**Screenshots**
N.A.

**GaNDLF Version**
0.0.10-dev

**Additional context**
N.A.



Discussion:
Getting this error:

```python-traceback
Exception has occurred: RuntimeError
running_mean should contain 64 elements not 32
  File "C:\Projects\GANDLF\GANDLF\models\seg_modules\DecodingModule.py", line 48, in forward
    x = self.act(self.in_0(x))
  File "C:\Projects\GANDLF\GANDLF\models\unet.py", line 335, in forward
    x = self.de_3(x, x4)
  File "C:\Projects\GANDLF\GANDLF\training_loop.py", line 84, in step
    output = model(image)
  File "C:\Projects\GANDLF\GANDLF\training_loop.py", line 175, in train_network
    loss, calculated_metrics, _ = step(model, image, label, params)
  File "C:\Projects\GANDLF\GANDLF\training_loop.py", line 810, in training_loop
    model, train_dataloader, optimizer, params
  File "C:\Projects\GANDLF\GANDLF\training_manager.py", line 242, in TrainingManager
    testing_data=testingData,
  File "C:\Projects\GANDLF\testing\test_full.py", line 157, in test_train_segmentation_rad_2d
    reset_prev=True,
  File "C:\Projects\GANDLF\gandlf_run", line 86, in main
    test_train_segmentation_rad_2d("cuda")
  File "C:\Projects\GANDLF\gandlf_run", line 158, in <module>
    main()
```
This is commented, for now.

Resolving patch:
diff --git a/GANDLF/models/__init__.py b/GANDLF/models/__init__.py
--- a/GANDLF/models/__init__.py
+++ b/GANDLF/models/__init__.py
@@ -1,4 +1,5 @@
 from .unet import unet, resunet
+from .light_unet import light_unet, light_resunet
 from .uinc import uinc
 from .fcn import fcn
 from .vgg import vgg11, vgg13, vgg16, vgg19
@@ -12,6 +13,10 @@
     "unet": unet,
     "resunet": resunet,
     "residualunet": unet,
+    "lightunet": light_unet,
+    "light_unet": light_unet,
+    "lightresunet": light_resunet,
+    "light_resunet": light_resunet,
     "fcn": fcn,
     "uinc": uinc,
     "vgg": vgg19,
diff --git a/GANDLF/models/light_unet.py b/GANDLF/models/light_unet.py
new file mode 100644
--- /dev/null
+++ b/GANDLF/models/light_unet.py
@@ -0,0 +1,192 @@
+# -*- coding: utf-8 -*-
+"""
+Implementation of UNet
+"""
+
+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 checkPatchDivisibility
+
+
+class light_unet(ModelBase):
+    """
+    This is the LIGHT U-Net architecture.
+    """
+
+    def __init__(
+        self,
+        parameters: dict,
+        residualConnections=False,
+    ):
+        self.network_kwargs = {"res": residualConnections}
+        super(light_unet, self).__init__(parameters)
+
+        self.network_kwargs = {"res": False}
+        
+        if not (checkPatchDivisibility(parameters["patch_size"])):
+            sys.exit(
+                "The patch size is not divisible by 16, which is required for",
+                parameters["model"]["architecture"],
+            )
+
+        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_0 = DownsamplingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            norm=self.Norm,
+        )
+        self.en_1 = EncodingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            dropout=self.Dropout,
+            norm=self.Norm,
+            network_kwargs=self.network_kwargs,
+        )
+        self.ds_1 = DownsamplingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            norm=self.Norm,
+        )
+        self.en_2 = EncodingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            dropout=self.Dropout,
+            norm=self.Norm,
+            network_kwargs=self.network_kwargs,
+        )
+        self.ds_2 = DownsamplingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            norm=self.Norm,
+        )
+        self.en_3 = EncodingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            dropout=self.Dropout,
+            norm=self.Norm,
+            network_kwargs=self.network_kwargs,
+        )
+        self.ds_3 = DownsamplingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            norm=self.Norm,
+        )
+        self.en_4 = EncodingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            dropout=self.Dropout,
+            norm=self.Norm,
+            network_kwargs=self.network_kwargs,
+        )
+        self.us_3 = UpsamplingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+        )
+        self.de_3 = DecodingModule(
+            input_channels=self.base_filters * 2,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            norm=self.Norm,
+            network_kwargs=self.network_kwargs,
+        )
+        self.us_2 = UpsamplingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+        )
+        self.de_2 = DecodingModule(
+            input_channels=self.base_filters * 2,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            norm=self.Norm,
+            network_kwargs=self.network_kwargs,
+        )
+        self.us_1 = UpsamplingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+        )
+        self.de_1 = DecodingModule(
+            input_channels=self.base_filters * 2,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+            norm=self.Norm,
+            network_kwargs=self.network_kwargs,
+        )
+        self.us_0 = UpsamplingModule(
+            input_channels=self.base_filters,
+            output_channels=self.base_filters,
+            conv=self.Conv,
+        )
+        self.out = out_conv(
+            input_channels=self.base_filters * 2,
+            output_channels=self.n_classes,
+            conv=self.Conv,
+            norm=self.Norm,
+            final_convolution_layer=self.final_convolution_layer,
+        )
+
+    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].
+
+        """
+        x1 = self.ins(x)
+        x2 = self.ds_0(x1)
+        x2 = self.en_1(x2)
+        x3 = self.ds_1(x2)
+        x3 = self.en_2(x3)
+        x4 = self.ds_2(x3)
+        x4 = self.en_3(x4)
+        x5 = self.ds_3(x4)
+        x5 = self.en_4(x5)
+
+        x = self.us_3(x5)
+        x = self.de_3(x, x4)
+        x = self.us_2(x)
+        x = self.de_2(x, x3)
+        x = self.us_1(x)
+        x = self.de_1(x, x2)
+        x = self.us_0(x)
+        x = self.out(x, x1)
+        return x
+
+
+class light_resunet(light_unet):
+    """
+    This is the LIGHT U-Net architecture with residual connections.
+    """
+
+    def __init__(self, parameters: dict):
+        super(light_resunet, self).__init__(parameters, residualConnections=True)
+
+
diff --git a/GANDLF/models/unet.py b/GANDLF/models/unet.py
--- a/GANDLF/models/unet.py
+++ b/GANDLF/models/unet.py
@@ -193,174 +193,3 @@ class resunet(unet):
     def __init__(self, parameters: dict):
         super(resunet, self).__init__(parameters, residualConnections=True)
 
-
-# class light_unet(ModelBase):
-#     """
-#     This is the standard U-Net architecture : https://arxiv.org/pdf/1606.06650.pdf. The 'residualConnections' flag controls residual connections The Downsampling, Encoding, Decoding 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,
-#         n_dimensions,
-#         n_channels,
-#         n_classes,
-#         base_filters,
-#         norm_type,
-#         final_convolution_layer,
-#         residualConnections=False,
-#     ):
-#         self.network_kwargs = {"res": False}
-#         super(light_unet, self).__init__(
-#             n_dimensions,
-#             n_channels,
-#             n_classes,
-#             base_filters,
-#             norm_type,
-#             final_convolution_layer,
-#         )
-#         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_0 = DownsamplingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             norm=self.Norm,
-#         )
-#         self.en_1 = EncodingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             dropout=self.Dropout,
-#             norm=self.Norm,
-#             network_kwargs=self.network_kwargs,
-#         )
-#         self.ds_1 = DownsamplingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             norm=self.Norm,
-#         )
-#         self.en_2 = EncodingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             dropout=self.Dropout,
-#             norm=self.Norm,
-#             network_kwargs=self.network_kwargs,
-#         )
-#         self.ds_2 = DownsamplingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             norm=self.Norm,
-#         )
-#         self.en_3 = EncodingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             dropout=self.Dropout,
-#             norm=self.Norm,
-#             network_kwargs=self.network_kwargs,
-#         )
-#         self.ds_3 = DownsamplingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             norm=self.Norm,
-#         )
-#         self.en_4 = EncodingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             dropout=self.Dropout,
-#             norm=self.Norm,
-#             network_kwargs=self.network_kwargs,
-#         )
-#         self.us_3 = UpsamplingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#         )
-#         self.de_3 = DecodingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             norm=self.Norm,
-#             network_kwargs=self.network_kwargs,
-#         )
-#         self.us_2 = UpsamplingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#         )
-#         self.de_2 = DecodingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             norm=self.Norm,
-#             network_kwargs=self.network_kwargs,
-#         )
-#         self.us_1 = UpsamplingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#         )
-#         self.de_1 = DecodingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#             norm=self.Norm,
-#             network_kwargs=self.network_kwargs,
-#         )
-#         self.us_0 = UpsamplingModule(
-#             input_channels=self.base_filters,
-#             output_channels=self.base_filters,
-#             conv=self.Conv,
-#         )
-#         self.out = out_conv(
-#             input_channels=self.base_filters,
-#             output_channels=n_classes,
-#             conv=self.Conv,
-#             norm=self.Norm,
-#             final_convolution_layer=self.final_convolution_layer,
-#         )
-
-#     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].
-
-#         """
-#         x1 = self.ins(x)
-#         x2 = self.ds_0(x1)
-#         x2 = self.en_1(x2)
-#         x3 = self.ds_1(x2)
-#         x3 = self.en_2(x3)
-#         x4 = self.ds_2(x3)
-#         x4 = self.en_3(x4)
-#         x5 = self.ds_3(x4)
-#         x5 = self.en_4(x5)
-
-#         x = self.us_3(x5)
-#         x = self.de_3(x, x4)
-#         x = self.us_2(x)
-#         x = self.de_2(x, x3)
-#         x = self.us_1(x)
-#         x = self.de_1(x, x2)
-#         x = self.us_0(x)
-#         x = self.out(x, x1)
-#         return x


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -20,6 +20,8 @@
 ## global defines
 # pre-defined segmentation model types for testing
 all_models_segmentation = [
+    "lightunet",
+    "lightresunet",
     "unet",
     "fcn",
     "uinc",