feat: residual conv blocks

src/neuron_proofreader/geometric_learning/curve_transformer.py

@@ -8,7 +8,6 @@
 
 """
 
-import numpy as np
 import torch
 import torch.nn as nn
 

src/neuron_proofreader/models/new_vision_models.py

@@ -31,6 +31,7 @@ def __init__(
         num_single_blocks=2,
         output_dim=1,
         use_double=True,
+        use_resblock=False,
     ):
         """
         Instantiates a CNN3D object.
@@ -57,9 +58,10 @@ def __init__(
             base_channels,
             depth,
             channel_multiplier=channel_multiplier,
-            use_double=use_double,
             max_channels=max_channels,
             num_single_blocks=num_single_blocks,
+            use_double=use_double,
+            use_resblock=use_resblock
         )
         self.output = FeedForwardNet(self.encode.out_channels, output_dim, 3)
 
@@ -110,17 +112,18 @@ def forward(self, x):
 
 class Encoder3D(nn.Module):
     """
-    Sequence of ConvBlock3D blocks with growing (capped) channel width.
+    Sequence of convolution blocks with growing (capped) channel width.
     """
     def __init__(
         self,
         in_channels,
         out_channels,
         depth,
         channel_multiplier=2,
-        use_double=True,
         max_channels=256,
         num_single_blocks=2,
+        use_double=True,
+        use_resblock=False,
     ):
         """
         Instantiates an Encoder3D object.
@@ -136,7 +139,7 @@ def __init__(
         channel_multiplier : float, optional
             Multiplicative channel growth factor per layer. Default is 2.
         use_double : bool, optional
-            Indication of whether blocks use double convolution. Default is True.
+            True if blocks use double convolution. Default is True.
         max_channels : int, optional
             Cap on channel growth across layers. Default is 128.
         """
@@ -145,9 +148,10 @@ def __init__(
 
         # Create encoding blocks
         blocks = list()
+        Block3D = ResConvBlock3D if use_resblock else ConvBlock3D
         for i in range(depth):
             use_double = i > num_single_blocks
-            block = ConvBlock3D(
+            block = Block3D(
                 in_channels, out_channels, use_double=use_double
             )
             blocks.append(block)
@@ -224,3 +228,84 @@ def get_num_groups(num_channels, max_groups=8):
             if num_channels % g == 0:
                 return g
         return 1
+
+
+class ResConvBlock3D(nn.Module):
+    """
+    3D convolutional block with a residual connection around the conv unit(s),
+    followed by downsampling. The skip path projects channels via a 1x1 conv
+    when in_channels != out_channels (or is an identity otherwise), so the
+    residual add is always shape-compatible.
+    """
+
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, use_double=True
+    ):
+        """
+        Instantiates a ResConvBlock3D object.
+
+        Parameters
+        ----------
+        in_channels : int
+            Number of input channels to this block.
+        out_channels : int
+            Number of output channels from this block.
+        kernel_size : int, optional
+            Size of kernel used on convolutional layers. Default is 3.
+        use_double : bool, optional
+            Indication of whether to apply a second conv+norm before the
+            residual add. Default is True.
+        """
+        # Call parent class
+        super().__init__()
+
+        # Instance attributes
+        self.out_channels = out_channels
+
+        # Architecture
+        layers = self.create_unit(in_channels, out_channels, kernel_size)
+        if use_double:
+            layers.extend(
+                self.create_unit(
+                    out_channels, out_channels, kernel_size, act=False
+                )
+            )
+        self.main = nn.Sequential(*layers)
+
+        # Skip path
+        if in_channels != out_channels:
+            self.skip = nn.Conv3d(in_channels, out_channels, kernel_size=1)
+        else:
+            self.skip = nn.Identity()
+
+        self.act = nn.GELU()
+        self.pool = nn.MaxPool3d(kernel_size=2)
+
+    def forward(self, x):
+        out = self.main(x) + self.skip(x)
+        out = self.act(out)
+        return self.pool(out)
+
+    # --- Helpers ---
+    def create_unit(self, in_channels, out_channels, kernel_size, act=True):
+        padding = kernel_size // 2
+        n_groups = self.get_num_groups(out_channels, 8)
+        unit = [
+            nn.Conv3d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                padding=padding,
+            ),
+            nn.GroupNorm(n_groups, out_channels),
+        ]
+        if act:
+            unit.append(nn.GELU())
+        return unit
+
+    @staticmethod
+    def get_num_groups(num_channels, max_groups=8):
+        for g in reversed(range(1, max_groups + 1)):
+            if num_channels % g == 0:
+                return g
+        return 1

src/neuron_proofreader/utils/ml_util.py

@@ -144,6 +144,8 @@ def update(self, pred, y, loss):
         self.loss += loss.item()
         self.n += y.numel()
 
+        del pred, y, loss
+
     def compute(self):
         precision = self.tp / (self.tp + self.fp + 1e-8)
         recall = self.tp / (self.tp + self.fn + 1e-8)