Dynamic UNet

Unet model using PixelShuffle ICNR upsampling that can be built on top of any pretrained architecture

UnetBlock

 UnetBlock (up_in_c, x_in_c, hook, final_div=True, blur=False,
            act_cls=<class 'torch.nn.modules.activation.ReLU'>,
            self_attention=False, init=<function 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', device=None,
            dtype=None)

A quasi-UNet block, using PixelShuffle_ICNR upsampling.

source

ResizeToOrig

 ResizeToOrig (mode='nearest')

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

source

DynamicUnet

 DynamicUnet (encoder, n_out, img_size, blur=False, blur_final=True,
              self_attention=False, y_range=None, last_cross=True,
              bottle=False, act_cls=<class
              'torch.nn.modules.activation.ReLU'>, init=<function
              kaiming_normal_>, norm_type=None, **kwargs)

Create a U-Net from a given architecture.

from fastai.vision.models import resnet34

m = resnet34()
m = nn.Sequential(*list(m.children())[:-2])
tst = DynamicUnet(m, 5, (128,128), norm_type=None)
x = cast(torch.randn(2, 3, 128, 128), TensorImage)
y = tst(x)
test_eq(y.shape, [2, 5, 128, 128])

tst = DynamicUnet(m, 5, (128,128), norm_type=None)
x = torch.randn(2, 3, 127, 128)
y = tst(x)