Dynamic UNet

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

/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0

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_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.

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)

©2021 fast.ai. All rights reserved.
Site last generated: Mar 31, 2021