| """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)] |
| 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() |
|
|
| 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) |
|
|
