body_3d = create_body(r3d_18, pretrained = False)

Dynamic Unet 3D

Fastai's DynamicUnet allows construction of a UNet using any pretrained CNN as backbone/encoder. A key module is nn.PixelShuffle which allows subpixel convolutions for upscaling in the UNet Blocks. However, nn.PixelShuffle is only for 2D images, so in faimed3d nn.ConvTranspose3d is used instead.

class ConvTranspose3D[source]

ConvTranspose3D(ni, nf=None, scale=2, blur=False, act_cls=None, norm_type=None, **kwargs) :: Sequential

Upsample by 2` from `ni` filters to `nf` (default `ni`), using `nn.ConvTranspose3D`.

Fastai's PixelShuffle_ICNR first performes a convolution to increase the layer size, then applies PixelShuffle to resize the image. A special initialization technique is applied to PixelShuffle, which can reduce checkerboard artifacts (see https://arxiv.org/pdf/1707.02937.pdf). It is probably not needed for nn.ConvTranspose3d

ConvTranspose3D(256, 128)(torch.randn((1, 256, 3, 13, 13))).size()

torch.Size([1, 128, 5, 15, 15])

ConvTranspose3D(256, 128, blur = True)(torch.randn((1, 256, 3, 13, 13))).size()

torch.Size([1, 128, 5, 15, 15])

To work with 3D data, the UnetBlock from fastai is adapted, replacing PixelShuffle_ICNR with the above created ConvTranspose3D and also adapting all conv-layers and norm-layers to the 3rd dimension. As small differences in size may appear, forward-func contains a interpolation step, which is also adapted to work with 5D input instead of 4D. UnetBlock3D receives the lower level features as hooks.

class UnetBlock3D[source]

UnetBlock3D(up_in_c, x_in_c, hook, final_div=True, blur=False, act_cls=ReLU, self_attention=False, init=kaiming_normal_, norm_type=None, ks=3, stride=1, padding=None, bias=None, ndim=2, bn_1st=True, transpose=False, xtra=None, bias_std=0.01, dilation:Union[int, Tuple[int, int]]=1, groups:int=1, padding_mode:str='zeros') :: Module

A quasi-UNet block, using `ConvTranspose3d` for upsampling`.

The output size of the last Unet-Block can be slightly different than the original input size, so one of the last steps in DynamicUnet is ResizeToOrig which is also adapted to work with 5D instead of 4D input images.

class ResizeToOrig[source]

ResizeToOrig(mode='nearest') :: Module

Merge a shortcut with the result of the module by adding them or concatenating them if `dense=True`.

DynamicUnet3D is the main UNet class for faimed3d and is very similar to the fastai DynamicUnet.

class DynamicUnet3D[source]

DynamicUnet3D(encoder, n_out, img_size, blur=False, blur_final=True, self_attention=False, y_range=None, last_cross=True, bottle=False, act_cls=ReLU, init=kaiming_normal_, norm_type=None, **kwargs) :: SequentialEx

Create a U-Net from a given architecture.

unet3d = DynamicUnet3D(body_3d, 2, (10,50,50), 1)

unet3d(torch.rand(2, 3, 10, 50, 50)).size()

torch.Size([2, 2, 10, 50, 50])