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

class UnetBlock[source]

UnetBlock(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 PixelShuffle_ICNR upsampling.

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.

class DynamicUnet[source]

DynamicUnet(encoder, n_classes, 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.

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 = torch.randn(2, 3, 128, 128)
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)