Added some unit tests for new functionalities

Author: StijnvWijn

tests/test_patch_gan.py

@@ -58,6 +58,65 @@
     [(1, 1, 32, 64, 32), (1, 1, 4, 8, 4)],
     [4, 7],
 ]
+TEST_3D_POOL = [
+{
+        "num_d": 2,
+        "num_layers_d": 3,
+        "spatial_dims": 3,
+        "num_channels": 8,
+        "in_channels": 3,
+        "out_channels": 1,
+        "kernel_size": 3,
+        "pooling_method": "max",
+        "activation": "LEAKYRELU",
+        "norm": "instance",
+        "bias": False,
+        "dropout": 0.1,
+        "minimum_size_im": 256,
+    },
+    torch.rand([1, 3, 256, 512, 256]),
+    [(1, 1, 32, 64, 32), (1, 1, 16, 32, 16)],
+    [4, 4],
+]
+TEST_2D_POOL = [
+    {
+        "num_d": 4,
+        "num_layers_d": 3,
+        "spatial_dims": 2,
+        "num_channels": 8,
+        "in_channels": 3,
+        "out_channels": 1,
+        "kernel_size": 3,
+        "pooling_method": "avg",
+        "activation": "LEAKYRELU",
+        "norm": "instance",
+        "bias": False,
+        "dropout": 0.1,
+        "minimum_size_im": 256,
+    },
+    torch.rand([1, 3, 256, 512]),
+    [(1, 1, 32, 64), (1, 1, 16, 32), (1, 1, 8, 16), (1, 1, 4, 8)],
+    [4, 4, 4, 4],
+]
+TEST_LAYER_LIST = [
+    {
+        "num_d": 3,
+        "num_layers_d": [3,4,5],
+        "spatial_dims": 2,
+        "num_channels": 8,
+        "in_channels": 3,
+        "out_channels": 1,
+        "kernel_size": 3,
+        "activation": "LEAKYRELU",
+        "norm": "instance",
+        "bias": False,
+        "dropout": 0.1,
+        "minimum_size_im": 256,
+    },
+    torch.rand([1, 3, 256, 512]),
+    [(1, 1, 32, 64), (1, 1, 16, 32), (1, 1, 8, 16)],
+    [4, 5, 6],
+]
 TEST_TOO_SMALL_SIZE = [
     {
         "num_d": 2,
@@ -74,9 +133,24 @@
         "minimum_size_im": 256,
     }
 ]
+TEST_MISMATCHED_NUM_LAYERS = [
+    {
+        "num_d": 5,
+        "num_layers_d": [3,4,5],
+        "spatial_dims": 2,
+        "num_channels": 8,
+        "in_channels": 3,
+        "out_channels": 1,
+        "kernel_size": 3,
+        "activation": "LEAKYRELU",
+        "norm": "instance",
+        "bias": False,
+        "dropout": 0.1,
+        "minimum_size_im": 256,
+    }
+]
 
-CASES = [TEST_2D, TEST_3D]
-
+CASES = [TEST_2D, TEST_3D, TEST_3D_POOL, TEST_2D_POOL, TEST_LAYER_LIST]
 
 class TestPatchGAN(unittest.TestCase):
     @parameterized.expand(CASES)
@@ -93,6 +167,10 @@ def test_too_small_shape(self):
         with self.assertRaises(AssertionError):
             MultiScalePatchDiscriminator(**TEST_TOO_SMALL_SIZE[0])
 
+    def test_mismatched_num_layers(self):
+        with self.assertRaises(AssertionError):
+            MultiScalePatchDiscriminator(**TEST_MISMATCHED_NUM_LAYERS[0])
+
     def test_script(self):
         net = MultiScalePatchDiscriminator(
             num_d=2,