--- 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') :: 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=False, dataset='wue1_cFOS', progress=True)

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=False, dataset='wue1_cFOS', progress=True)

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=False, dataset='wue_cFOS', progress=True, dropout=0.5)

U-Net model optimized for deepflash2

{% endraw %} {% raw %}
{% endraw %} {% raw %}
tst = unet_deepflash2()
tst = unet_deepflash2(pretrained=True)
tst = unet_deepflash2(n_classes=3, pretrained=True)
No pretrained weights for 3 classes in final layer.
{% endraw %} {% raw %}

unet_custom[source]

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

Customizable U-Net model. Customize via kwargs

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