feat: unet architecture from n2v2#30
Conversation
carshadi
left a comment
carshadi
left a comment
There was a problem hiding this comment.
Hi @anna-grim , this is a valuable idea. As-is, this code will raise an error with the default values so please take a look at the comments. It would be useful to add some unit tests to verify that these issues are resolved. In particular, a test exercising remove_top_skip=True with odd spatial dimensions and the UNet() constructor issue.
| if ( | ||
| isinstance(width_multiplier, bool) | ||
| or not isinstance(width_multiplier, Real) | ||
| isinstance(width_multiplier, Real) |
There was a problem hiding this comment.
Why change this condition? isinstance(width_multiplier, Real) is True for every ordinary int or float so attempting to construct this class will always raise a ValueError.
| scale_factor=2, mode="trilinear", align_corners=True | ||
| ) | ||
| self.conv = DoubleConv( | ||
| in_channels, out_channels, mid_channels=in_channels // 2 |
There was a problem hiding this comment.
When loading a checkpoint created before this patch with trilinear=True, omitting
mid_channels=in_channels // 2 changes parameter shapes throughout up1–up3. Thus, load_state_dict in load_model will reject existing checkpoints. This architecture should probably be a separate class, e.g., UNetN2V2 so we can preserve backwards compatibility and do A/B testing more easily.
|
|
||
| x = torch.cat([x2, x1], dim=1) | ||
|
|
||
| else: |
There was a problem hiding this comment.
When remove_top_skip=True and any input spatial dimension is odd, the final upsampling doubles floor(input_size / 2) without aligning it to x1, leaving the output one voxel smaller in that dimension. With the default residual mode, x + logits then raises a size-mismatch error; without residual mode, the output tensor still has the wrong shape.
Example:
Using an odd-sized input (1, 1, 17, 19, 21) with remove_top_skip=True:
residual=False:
input = (1, 1, 17, 19, 21)
output = (1, 1, 16, 18, 20)
residual=True:
RuntimeError: The size of tensor a (21) must match the size of tensor b (20)
The encoder’s first pooling produces (8, 9, 10). The final decoder upsampling doubles that to (16, 18, 20), losing one voxel along each odd dimension.
| trilinear=True, | ||
| residual=True, | ||
| maxblurpool=False, | ||
| remove_top_skip=False, |
There was a problem hiding this comment.
The default residual=True would defeat the purpose of remove_top_skip=True. If a user enables remove_top_skip=True to reproduce N2V2's hot-pixel/checkerboard suppression but leaves residual at its default True; pixel noise still passes straight through via the residual add, so the architectural change does nothing and the experiment silently evaluates the wrong hypothesis. At minimum document the interaction; consider warning or disallowing remove_top_skip=True with residual=True.
| x, | ||
| self.kernel, | ||
| stride=2, | ||
| padding=1, |
There was a problem hiding this comment.
MaxBlurPool3D's blur conv uses zero padding, so the normalized [1,2,1] binomial kernel systematically attenuates border voxels at every encoder scale (reference BlurPool implementations use reflect/replicate padding e.g., https://github.com/seominseok0429/Impact-of-Aliasing-on-Generalization-in-Deep-Convolutional-Networks-pytorch/blob/main/blurpool.py). For example, constant input 1.0 yields interior 1.0 but faces 0.75, edges 0.56, corners 0.42, vs 1.0 everywhere for the nn.MaxPool3d(2) it replaces.
2 participants
No description provided.