From 580552b2433c98e4310b90bb4f0bedfb28debd11 Mon Sep 17 00:00:00 2001 From: anna-grim Date: Sat, 11 Jul 2026 01:14:52 +0000 Subject: [PATCH] feat: residual conv blocks --- .../geometric_learning/curve_transformer.py | 1 - .../models/new_vision_models.py | 95 ++++++++++++++++++- src/neuron_proofreader/utils/ml_util.py | 2 + 3 files changed, 92 insertions(+), 6 deletions(-) diff --git a/src/neuron_proofreader/geometric_learning/curve_transformer.py b/src/neuron_proofreader/geometric_learning/curve_transformer.py index 0e5b667..676ecc0 100644 --- a/src/neuron_proofreader/geometric_learning/curve_transformer.py +++ b/src/neuron_proofreader/geometric_learning/curve_transformer.py @@ -8,7 +8,6 @@ """ -import numpy as np import torch import torch.nn as nn diff --git a/src/neuron_proofreader/models/new_vision_models.py b/src/neuron_proofreader/models/new_vision_models.py index 5548e9e..78327e5 100644 --- a/src/neuron_proofreader/models/new_vision_models.py +++ b/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,7 +112,7 @@ 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, @@ -118,9 +120,10 @@ def __init__( 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 diff --git a/src/neuron_proofreader/utils/ml_util.py b/src/neuron_proofreader/utils/ml_util.py index fa42897..562aa6e 100644 --- a/src/neuron_proofreader/utils/ml_util.py +++ b/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)