vision.models.unet
Dynamic Unet that can use any pretrained model as a backbone.
Dynamic U-Net
This module builds a dynamic U-Net from any backbone pretrained on ImageNet, automatically inferring the intermediate sizes.
This is the original U-Net. The difference here is that the left part is a pretrained model.
class
DynamicUnet
[source][test]
DynamicUnet
(encoder
:Module
,n_classes
:int
,img_size
:Tuple
[int
,int
]=(256, 256)
,blur
:bool
=False
,blur_final
=True
,self_attention
:bool
=False
,y_range
:OptRange
=None
,last_cross
:bool
=True
,bottle
:bool
=False
, **kwargs
) ::PrePostInitMeta
::SequentialEx
Tests found forDynamicUnet
:
pytest -sv tests/test_vision_models_unet.py::test_dynamic_unet_shape
[source]pytest -sv tests/test_vision_models_unet.py::test_unet_block_shapes
[source]
To run tests please refer to this guide.
Create a U-Net from a given architecture.
This U-Net will sit on top of an encoder
(that can be a pretrained model) and with a final output of n_classes
. During the initialization, it uses Hooks
to determine the intermediate features sizes by passing a dummy input through the model and create the upward path automatically.
blur
is used to avoid checkerboard artifacts at each layer, blur_final
is specific to the last layer. self_attention
determines if we use a self attention layer at the third block before the end. If y_range
is passed, the last activations go through a sigmoid rescaled to that range. last_cross
determines if we use a cross-connection with the direct input of the model, and in this case bottle
flags if we use a bottleneck or not for that skip connection.
class
UnetBlock
[source][test]
UnetBlock
(up_in_c
:int
,x_in_c
:int
,hook
:Hook
,final_div
:bool
=True
,blur
:bool
=False
,leaky
:float
=None
,self_attention
:bool
=False
, **kwargs
) ::PrePostInitMeta
::Module
No tests found forUnetBlock
. To contribute a test please refer to this guide and this discussion.
A quasi-UNet block, using PixelShuffle_ICNR upsampling
.
This block receives the output of the last block to be upsampled (size up_in_c
) and the activations features from an intermediate layer of the encoder
(size x_in_c
, this is the lateral connection). The hook
is set to this intermediate layer to store the output needed for this block. final_div
determines if we divide the number of features by 2 with the upsampling, blur
is to avoid checkerboard artifacts. If leaky
is set, use a leaky ReLU with a slope equals to that parameter instead of a ReLU, and self_attention
determines if we use a self-attention layer. kwargs
are passed to conv_layer
.
©2021 fast.ai. All rights reserved.
Site last generated: Jan 5, 2021