--- title: Models keywords: fastai sidebar: home_sidebar summary: "Pytorch segmentation models." description: "Pytorch segmentation models." nb_path: "nbs/01_models.ipynb" ---
{% raw %}
{% endraw %} {% raw %}
{% endraw %}

U-Net models

Pytorch implementation adapted from https://github.com/jvanvugt/pytorch-unet

{% raw %}

class UNetConvBlock[source]

UNetConvBlock(in_size, out_size, padding, batch_norm, dropout=0.0, neg_slope=0.1) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

{% endraw %} {% raw %}
{% endraw %} {% raw %}

class UNetUpBlock[source]

UNetUpBlock(in_size, out_size, up_mode, padding, batch_norm, dropout=0.0, neg_slope=0.1) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

{% endraw %} {% raw %}
{% endraw %} {% raw %}

class UNet2D[source]

UNet2D(in_channels=1, n_classes=2, depth=5, wf=6, padding=False, batch_norm=False, dropout=0.0, neg_slope=0.0, up_mode='upconv', **kwargs) :: Module

Pytorch U-Net Implementation

{% endraw %} {% raw %}
{% endraw %}

Args:\ in_channels (int): the number of input channels.\ n_classes (int): the number of output channels. \ depth (int): depth of the network.\ wf (int): number of filters in the first layer is 2^wf padding (bool): if True, apply padding such that the input shape is the same as the output. This may introduce artifacts\ batch_norm (bool): Use BatchNorm after layers with an activation function\ up_mode (str): one of 'upconv' or 'upsample'. 'upconv' will use transposed convolutions for learned upsampling. 'upsample' will use bilinear upsampling.\ neg_slope(float): Controls the angle of the negative slope for LeakyReLU. Standard ReLU if set to 0.

{% raw %}
{% endraw %}

U-Net architectures

Original U-Net architecture based on Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. "U-net: Convolutional networks for biomedical image segmentation." International Conference on Medical image computing and computer-assisted intervention. Springer, Cham, 2015.

{% raw %}

unet_ronneberger2015[source]

unet_ronneberger2015(in_channels=1, n_classes=2, pretrained=None, progress=True, **kwargs)

Original U-Net architecture based on Ronnberger et al. (2015)

{% endraw %} {% raw %}
{% endraw %}

U-Net architecture based on Falk, Thorsten, et al. "U-Net: deep learning for cell counting, detection, and morphometry." Nature methods 16.1 (2019): 67-70.

{% raw %}

unet_falk2019[source]

unet_falk2019(in_channels=1, n_classes=2, pretrained=None, progress=True, **kwargs)

U-Net architecture based on Falk et al. (2019)

{% endraw %} {% raw %}
{% endraw %}

U-Net model optimized for deepflash2

{% raw %}

unet_deepflash2[source]

unet_deepflash2(in_channels=1, n_classes=2, pretrained=None, progress=True, dropout=0.5, **kwargs)

U-Net model optimized for deepflash2

{% endraw %} {% raw %}
{% endraw %} {% raw %}
tst = unet_deepflash2()
tst = unet_deepflash2(pretrained='wue_cFOS')
tst = unet_deepflash2(n_classes=3, pretrained='wue_cFOS')
Loaded model weights trained on wue_cFOS.
No pretrained weights for 3 classes in final layer.
Loaded model weights trained on wue_cFOS.
{% endraw %} {% raw %}

unet_custom[source]

unet_custom(in_channels=1, n_classes=2, pretrained=None, progress=True, **kwargs)

Customizable U-Net model. Customize via kwargs

{% endraw %} {% raw %}
{% endraw %}

UneXt

UneXt50 Architecture adapted from Maxim Shugaev on Kaggle

{% raw %}

class FPN[source]

FPN(input_channels:list, output_channels:list) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

{% endraw %} {% raw %}
{% endraw %} {% raw %}

class UnetBlock[source]

UnetBlock(up_in_c:int, x_in_c:int, nf:int=None, blur:bool=False, self_attention:bool=False, padding:int=1, **kwargs) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

{% endraw %} {% raw %}
{% endraw %} {% raw %}
{% endraw %} {% raw %}

class ASPP[source]

ASPP(inplanes=512, mid_c=256, dilations=[6, 12, 18, 24], out_c=None) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

{% endraw %} {% raw %}
{% endraw %} {% raw %}

class UneXt50[source]

UneXt50(in_channels=1, n_classes=2, stride=1, inplanes=64, pre_ssl=True, **kwargs) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

{% endraw %} {% raw %}
{% endraw %} {% raw %}

unext50_deepflash2[source]

unext50_deepflash2(in_channels=1, n_classes=2, pretrained=None, progress=True, **kwargs)

UneXt50 model. Customize via kwargs

{% endraw %} {% raw %}
{% endraw %} {% raw %}
tst = unext50_deepflash2(pretrained='wue_cFOS')
tst = unext50_deepflash2(in_channels=3, n_classes=5)
x = torch.randn(2, 3, 518, 518)
y = tst(x)
test_eq(y.shape, [2, 5, 392, 392])
Using cache found in /media/data/home/mag01ud/.cache/torch/hub/facebookresearch_semi-supervised-ImageNet1K-models_master
Loaded model weights trained on wue_cFOS.
Using cache found in /media/data/home/mag01ud/.cache/torch/hub/facebookresearch_semi-supervised-ImageNet1K-models_master
{% endraw %}

Shape Defaults

Helper functions for default input and masks shapes, depending on model architecture

{% raw %}

get_default_shapes[source]

get_default_shapes(arch)

{% endraw %} {% raw %}
{% endraw %}