File size: 4,392 Bytes
c0aeb30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Custom 3D residual U-Net (~10.7M params), valid-padding throughout.
Valid (unpadded) convolutions are used instead of same-padding because this setup is intended to run
under blockwise inference (daisy/gunpowder tiling adjacent crops across a large volume). Same-padding
convs synthesize zero-context at every block's border, which produces seam artifacts at block
boundaries; valid convs give every output voxel genuine surrounding context, matching the convention
already used by fly_organelles.model.StandardUnet and fish_trainer.model.Unet3d.
Verified (base_channels=40, factor=2, blocks_per_level=1): 10,702,493 params,
input (148,148,148) -> output (60,60,60).
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
def center_crop_to(x, target):
diffs = [x.shape[i] - target.shape[i] for i in (2, 3, 4)]
starts = [d // 2 for d in diffs]
return x[
:,
:,
starts[0] : starts[0] + target.shape[2],
starts[1] : starts[1] + target.shape[3],
starts[2] : starts[2] + target.shape[4],
]
class ResBlock(nn.Module):
"""Two valid 3x3x3 convs + InstanceNorm + ReLU, with a crop-and-add identity skip
(1x1x1 conv on the skip path when channel count changes; identity is center-cropped
to match the main path's spatially-shrunk output before the residual add)."""
def __init__(self, c_in, c_out):
super().__init__()
self.conv1 = nn.Conv3d(c_in, c_out, 3, padding=0)
self.norm1 = nn.InstanceNorm3d(c_out, affine=True)
self.conv2 = nn.Conv3d(c_out, c_out, 3, padding=0)
self.norm2 = nn.InstanceNorm3d(c_out, affine=True)
self.act = nn.ReLU(inplace=True)
self.skip = nn.Conv3d(c_in, c_out, 1) if c_in != c_out else nn.Identity()
def forward(self, x):
idn = self.skip(x)
h = self.act(self.norm1(self.conv1(x)))
h = self.norm2(self.conv2(h))
idn = center_crop_to(idn, h)
return self.act(h + idn)
class Down(nn.Module):
"""Exact 2x downsample (kernel==stride, divides evenly, no cropping needed)."""
def __init__(self, c_in, c_out):
super().__init__()
self.conv = nn.Conv3d(c_in, c_out, 2, stride=2)
def forward(self, x):
return self.conv(x)
class Up(nn.Module):
"""Non-learnable nearest-neighbor upsample (exact 2x) + 1x1x1 channel-reduce conv.
Deliberately non-learnable to match this codebase's own constant_upsample=True choice in
StandardUnet (avoids checkerboard/seam artifacts from learned transposed convs)."""
def __init__(self, c_in, c_out):
super().__init__()
self.reduce = nn.Conv3d(c_in, c_out, 1)
def forward(self, x):
x = F.interpolate(x, scale_factor=2, mode="nearest")
return self.reduce(x)
class ResUNet3DValid(nn.Module):
def __init__(self, out_channels, base_channels=40, factor=2, blocks_per_level=1, in_channels=1):
super().__init__()
C = [base_channels * (factor**i) for i in range(4)] # 3 down/up levels + bottleneck
self.C = C
self.stem = ResBlock(in_channels, C[0])
self.downs = nn.ModuleList()
self.enc_blocks = nn.ModuleList()
for i in range(3):
self.downs.append(Down(C[i], C[i + 1]))
self.enc_blocks.append(nn.ModuleList([ResBlock(C[i + 1], C[i + 1]) for _ in range(blocks_per_level)]))
self.ups = nn.ModuleList()
self.dec_blocks = nn.ModuleList()
for i in reversed(range(3)):
self.ups.append(Up(C[i + 1], C[i]))
blocks = [ResBlock(C[i] * 2, C[i])] + [ResBlock(C[i], C[i]) for _ in range(blocks_per_level - 1)]
self.dec_blocks.append(nn.ModuleList(blocks))
self.final_conv = nn.Conv3d(C[0], out_channels, 1)
def forward(self, input):
skips = []
x = self.stem(input)
skips.append(x)
for down, enc in zip(self.downs, self.enc_blocks):
x = down(x)
for blk in enc:
x = blk(x)
skips.append(x)
skips.pop() # bottleneck itself is not a skip target
for up, dec, skip in zip(self.ups, self.dec_blocks, reversed(skips)):
x = up(x)
skip_c = center_crop_to(skip, x)
x = torch.cat([x, skip_c], dim=1)
for blk in dec:
x = blk(x)
return self.final_conv(x)
|