feat: unet architecture from n2v2

src/aind_exaspim_image_compression/machine_learning/train.py

@@ -16,7 +16,6 @@
 import numpy as np
 import os
 import torch
-import torch.nn as nn
 import torch.optim as optim
 from skimage import io
 

src/aind_exaspim_image_compression/machine_learning/unet3d.py

@@ -38,7 +38,14 @@ class UNet(nn.Module):
         Final 1x1x1 convolution mapping features to the output channel.
     """
 
-    def __init__(self, width_multiplier=1, trilinear=True, residual=True):
+    def __init__(
+        self,
+        width_multiplier=1,
+        trilinear=True,
+        residual=True,
+        maxblurpool=False,
+        remove_top_skip=False,
+    ):
         """
         Instantiates a UNet object.
 
@@ -56,38 +63,87 @@ def __init__(self, width_multiplier=1, trilinear=True, residual=True):
             Default is True.
         """
         # Call parent class
-        super(UNet, self).__init__()
+        super().__init__()
 
         if (
-            isinstance(width_multiplier, bool)
-            or not isinstance(width_multiplier, Real)
+            isinstance(width_multiplier, Real)
             or width_multiplier < 1
             or not float(width_multiplier).is_integer()
         ):
             raise ValueError("width_multiplier must be a positive integer")
 
         # Initializations
-        _channels = (32, 64, 128, 256, 512)
+        base_channels = (32, 64, 128, 256, 512)
         factor = 2 if trilinear else 1
 
         # Instance attributes
         self.width_multiplier = int(width_multiplier)
-        self.channels = [c * self.width_multiplier for c in _channels]
+        self.channels = [
+            c * self.width_multiplier
+            for c in base_channels
+        ]
+
         self.trilinear = trilinear
         self.residual = residual
+        self.maxblurpool = maxblurpool
+        self.remove_top_skip = remove_top_skip
 
-        # Contracting layers
+        # Encoder
         self.inc = DoubleConv(1, self.channels[0])
-        self.down1 = Down(self.channels[0], self.channels[1])
-        self.down2 = Down(self.channels[1], self.channels[2])
-        self.down3 = Down(self.channels[2], self.channels[3])
-        self.down4 = Down(self.channels[3], self.channels[4] // factor)
-
-        # Expanding layers
-        self.up1 = Up(self.channels[4], self.channels[3] // factor, trilinear)
-        self.up2 = Up(self.channels[3], self.channels[2] // factor, trilinear)
-        self.up3 = Up(self.channels[2], self.channels[1] // factor, trilinear)
-        self.up4 = Up(self.channels[1], self.channels[0], trilinear)
+
+        self.down1 = Down(
+            self.channels[0],
+            self.channels[1],
+            maxblurpool=maxblurpool,
+        )
+
+        self.down2 = Down(
+            self.channels[1],
+            self.channels[2],
+            maxblurpool=maxblurpool,
+        )
+
+        self.down3 = Down(
+            self.channels[2],
+            self.channels[3],
+            maxblurpool=maxblurpool,
+        )
+
+        self.down4 = Down(
+            self.channels[3],
+            self.channels[4] // factor,
+            maxblurpool=maxblurpool,
+        )
+
+        # Decoder
+        self.up1 = Up(
+            self.channels[4],
+            self.channels[3] // factor,
+            trilinear=trilinear,
+            use_skip=True,
+        )
+
+        self.up2 = Up(
+            self.channels[3],
+            self.channels[2] // factor,
+            trilinear=trilinear,
+            use_skip=True,
+        )
+
+        self.up3 = Up(
+            self.channels[2],
+            self.channels[1] // factor,
+            trilinear=trilinear,
+            use_skip=True,
+        )
+
+        self.up4 = Up(
+            self.channels[1],
+            self.channels[0],
+            trilinear=trilinear,
+            use_skip=not remove_top_skip,
+        )
+
         self.outc = OutConv(self.channels[0], 1)
 
     @property
@@ -97,6 +153,8 @@ def config(self):
             "width_multiplier": self.width_multiplier,
             "trilinear": self.trilinear,
             "residual": self.residual,
+            "maxblurpool": self.maxblurpool,
+            "remove_top_skip": self.remove_top_skip,
         }
 
     def forward(self, x):
@@ -114,23 +172,29 @@ def forward(self, x):
             Output tensor with shape (B, 1, D, H, W), representing the
             denoised image.
         """
-        # Contracting layers
+        # Encoder
         x1 = self.inc(x)
         x2 = self.down1(x1)
         x3 = self.down2(x2)
         x4 = self.down3(x3)
         x5 = self.down4(x4)
 
-        # Expanding layers
+        # Decoder
         d = self.up1(x5, x4)
         d = self.up2(d, x3)
         d = self.up3(d, x2)
-        d = self.up4(d, x1)
+
+        if self.remove_top_skip:
+            d = self.up4(d)
+        else:
+            d = self.up4(d, x1)
+
         logits = self.outc(d)
 
         # Residual denoising: predict the correction added to the input
         if self.residual:
             return x + logits
+
         return logits
 
 
@@ -211,15 +275,14 @@ def forward(self, x):
 class Down(nn.Module):
     """
     A downsampling module for a 3D U-Net.
-
-    Attributes
-    ----------
-    maxpool_conv : nn.Sequential
-        Sequential module containing a MaxPool3d layer followed by a
-        DoubleConv block.
     """
 
-    def __init__(self, in_channels, out_channels):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        maxblurpool=False,
+    ):
         """
         Instantiates a Down object.
 
@@ -229,13 +292,23 @@ def __init__(self, in_channels, out_channels):
             Number of input channels to this module.
         out_channels : int
             Number of output channels produced by this module.
+        maxblurpool : bool, optional
+            True if max-blur pooling should be used to downsample. Default is
+            False.
         """
         # Call parent class
         super().__init__()
 
+        # Initializations
+        if maxblurpool:
+            downsample = MaxBlurPool3D(in_channels)
+        else:
+            downsample = nn.MaxPool3d(2)
+
         # Instance attributes
         self.maxpool_conv = nn.Sequential(
-            nn.MaxPool3d(2), DoubleConv(in_channels, out_channels)
+            downsample,
+            DoubleConv(in_channels, out_channels),
         )
 
     def forward(self, x):
@@ -269,7 +342,13 @@ class Up(nn.Module):
         connection.
     """
 
-    def __init__(self, in_channels, out_channels, trilinear=True):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        trilinear=True,
+        use_skip=True,
+    ):
         """
         Instantiates an Up object.
 
@@ -287,20 +366,32 @@ def __init__(self, in_channels, out_channels, trilinear=True):
         super().__init__()
 
         # Instance attributes
+        self.use_skip = use_skip
+
         if trilinear:
             self.up = nn.Upsample(
-                scale_factor=2, mode="trilinear", align_corners=True
-            )
-            self.conv = DoubleConv(
-                in_channels, out_channels, mid_channels=in_channels // 2
+                scale_factor=2,
+                mode="trilinear",
+                align_corners=True,
             )
         else:
             self.up = nn.ConvTranspose3d(
-                in_channels, in_channels // 2, kernel_size=2, stride=2
+                in_channels,
+                in_channels // 2,
+                kernel_size=2,
+                stride=2,
             )
-            self.conv = DoubleConv(in_channels, out_channels)
 
-    def forward(self, x1, x2):
+        conv_in_channels = (
+            in_channels if use_skip else in_channels // 2
+        )
+
+        self.conv = DoubleConv(
+            conv_in_channels,
+            out_channels,
+        )
+
+    def forward(self, x1, x2=None):
         """
         Forward pass of the upsampling block in a 3D U-Net.
 
@@ -309,7 +400,7 @@ def forward(self, x1, x2):
         x1 : torch.Tensor
             Input tensor from the previous decoder layer with shape
             (B, C1, D, H1, W1).
-        x2 : torch.Tensor
+        x2 : torch.Tensor, optional
             Skip connection tensor from the encoder path with shape
             (B, C2, D, H2, W2).
 
@@ -321,17 +412,65 @@ def forward(self, x1, x2):
             (B, out_channels, D, H2, W2).
         """
         x1 = self.up(x1)
-        diffY = x2.size()[2] - x1.size()[2]
-        diffX = x2.size()[3] - x1.size()[3]
 
-        x1 = F.pad(
-            x1,
-            [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2],
-        )
-        x = torch.cat([x2, x1], dim=1)
+        if self.use_skip:
+            diff_z = x2.size(2) - x1.size(2)
+            diff_y = x2.size(3) - x1.size(3)
+            diff_x = x2.size(4) - x1.size(4)
+
+            x1 = F.pad(
+                x1,
+                [
+                    diff_x // 2,
+                    diff_x - diff_x // 2,
+                    diff_y // 2,
+                    diff_y - diff_y // 2,
+                    diff_z // 2,
+                    diff_z - diff_z // 2,
+                ],
+            )
+
+            x = torch.cat([x2, x1], dim=1)
+
+        else:
+            x = x1
+
         return self.conv(x)
 
 
+class MaxBlurPool3D(nn.Module):
+
+    def __init__(self, channels):
+        super().__init__()
+
+        self.pool = nn.MaxPool3d(2, stride=1)
+
+        kernel = torch.tensor([1., 2., 1.])
+        kernel = (
+            kernel[:, None, None]
+            * kernel[None, :, None]
+            * kernel[None, None, :]
+        )
+        kernel /= kernel.sum()
+
+        self.register_buffer(
+            "kernel",
+            kernel[None, None].repeat(channels, 1, 1, 1, 1),
+        )
+        self.channels = channels
+
+    def forward(self, x):
+        x = self.pool(x)
+        x = F.conv3d(
+            x,
+            self.kernel,
+            stride=2,
+            padding=1,
+            groups=self.channels,
+        )
+        return x
+
+
 class OutConv(nn.Module):
     """
     Final output convolution layer for a 3D U-Net.