feat: GIN, a cascade of random convolutions for domain generalisation

README.md

@@ -232,6 +232,7 @@ column -- see "Run Monai training with SmaugLab augmentations" above. Regenerate
 | laplace | TA | `RandomLaplaceGPU` | `LaplaceConvTransform` |
 | unsharp_mask | TA | `RandomUnsharpMaskGPU` | — |
 | rand_conv | TA | `RandomRandConvGPU` | — |
+| gin | TA | `RandomGINGPU` | — |
 | bias_field | TA | `RandomBiasFieldGPU` | — |
 | inverse | TA | `RandomInverseGPU` | — |
 | histogram_equal | TA | `RandomHistogramEqualizationGPU` | `HistogramEqualTransform` |

smauglab/configs/all_augmentations.json

@@ -293,6 +293,32 @@
             "retain_stats": false,
             "same_on_batch": false
         },
+        "RandomGINGPU": {
+            "alpha_range": [
+                0.0,
+                1.0
+            ],
+            "apply_to_channel": [
+                0
+            ],
+            "in_seg": 0.0,
+            "interm_channels": 2,
+            "keepdim": true,
+            "kernel_sizes": [
+                1,
+                3
+            ],
+            "mix_in_out": false,
+            "n_layers": 4,
+            "negative_slope": 0.01,
+            "out_norm": "frob",
+            "out_seg": 0.0,
+            "p": 1.0,
+            "p_batch": 1.0,
+            "padding_mode": "reflect",
+            "retain_stats": false,
+            "same_on_batch": false
+        },
         "RandomClampGPU": {
             "apply_to_channel": [
                 0

smauglab/configs/transform_params_gpu.json

@@ -122,6 +122,9 @@
             "mix_prob": 0.5,
             "p": 0.2
         },
+        "RandomGINGPU": {
+            "p": 0.0
+        },
         "RandomClampGPU": {
             "max_clamp_amount": 0.2,
             "retain_stats": true,

smauglab/registry.py

@@ -113,6 +113,11 @@ class AugId(str, Enum):
     LAPLACE = "laplace"
     UNSHARP_MASK = "unsharp_mask"
     RAND_CONV = "rand_conv"
+    # A separate concept from RAND_CONV rather than a second implementation of it:
+    # RandConv is one random kernel, GIN is a cascade of them with a non-linearity
+    # between the layers and the output renormalised to the input's energy. Sharing the
+    # id would overwrite RandomRandConvGPU's matrix cell.
+    GIN = "gin"
     BIAS_FIELD = "bias_field"
     INVERSE = "inverse"
     HISTOGRAM_EQUAL = "histogram_equal"
@@ -213,6 +218,9 @@ def __init__(self, source: str, problems: list[str]) -> None:
             "RandomLaplaceGPU",
             "RandomUnsharpMaskGPU",  # UnsharpMaskTransform
             "RandomRandConvGPU",  # RandomConvTransform
+            # No ladder key: GIN postdates the ladder. Placed next to the single-layer
+            # RandConv it generalises, so a config reads the two together.
+            "RandomGINGPU",
             "RandomClampGPU",  # ClampTransform
             "RandomGaussianNoiseGPU",  # GaussianNoiseTransform
             "RandomGaussianBlurGPU",  # GaussianBlurTransform

smauglab/transforms/__init__.py

@@ -19,6 +19,7 @@ class carries an `@register(...)` decorator, so the registry is complete once th
 from smauglab.transforms.gpu import contrast as gpu_contrast  # noqa: F401
 from smauglab.transforms.gpu import domain_transfer  # noqa: F401
 from smauglab.transforms.gpu import fromSeg as gpu_fromSeg
+from smauglab.transforms.gpu import gin as gpu_gin  # noqa: F401
 from smauglab.transforms.gpu import palette as gpu_palette  # noqa: F401
 from smauglab.transforms.gpu import spatial as gpu_spatial  # noqa: F401
 from smauglab.transforms.synthseg import transforms as synthseg_transforms  # noqa: F401

unit_tests/test_transforms_gpu.py

@@ -25,6 +25,7 @@
     "smauglab.transforms.gpu.contrast",
     "smauglab.transforms.gpu.spatial",
     "smauglab.transforms.gpu.fromSeg",
+    "smauglab.transforms.gpu.gin",
     "smauglab.transforms.gpu.domain_transfer",
     "smauglab.transforms.gpu.palette.transform",
 ]