Creating and training a U-Net model with PyTorch for 2D & 3D semantic segmentation: Model building [2/4]
By: Johannes Schmidt
Tags: Unet, Python, Pytorch, Deep Learning, Semantic Segmentation

In the previous chapter we built a dataloader that picks up our images and performs some transformations and augmentations so that they can be fed in batches to a neural network like the U-Net. In this part, we focus on building a U-Net from scratch with the PyTorch library. You should be familiar with building a neural network with PyTorch and you should generally know the U-Net. The goal is to implement the U-Net in such a way, that important model configurations such as the activation function or the depth can be passed as arguments when creating the model.

About the U-Net

The U-Net is a convolutional neural network architecture that is designed for fast and precise segmentation of images. It has performed extremely well in several challenges and to this day, it is one of the most popular end-to-end architectures in the field of semantic segmentation.

We can split the network into two parts: The encoder path (backbone) and the decoder path. The encoder captures features at different scales of the images by using a traditional stack of convolutional and max pooling layers. Concretely speaking, a block in the encoder consists of the repeated use of two convolutional layers (k=3, s=1), each followed by a non-linearity layer, and a max-pooling layer (k=2, s=2). For every convolution block and its associated max pooling operation, the number of feature maps is doubled to ensure that the network can learn the complex structures effectively.

The decoder path is a symmetric expanding counterpart that uses transposed convolutions. This type of convolutional layer is an up-sampling method with trainable parameters and performs the reverse of (down)pooling layers such as the max pool. Similar to the encoder, each convolution block is followed by such an up-convolutional layer. The number of feature maps is halved in every block. Because recreating a segmentation mask from a small feature map is a rather difficult task for the network, the output after every up-convolutional layer is appended by the feature maps of the corresponding encoder block. The feature maps of the encoder layer are cropped if the dimensions exceed the one of the corresponding decoder layers.

In the end, the output passes another convolution layer (k=1, s=1) with the number of feature maps being equal to the number of defined labels. The result is a u-shaped convolutional network that offers an elegant solution for good localization and use of context. Let’s take a look at the code.

The code

This code is based on https://github.com/ELEKTRONN/elektronn3/blob/master/elektronn3/models/unet.py (c) 2017 Martin Drawitsch, released under MIT License, which implements a configurable (2D/3D) U-Net with user-defined network depth and a few other improvements of the original architecture. They themselves actually used the 2D code from Jackson Huang https://github.com/jaxony/unet-pytorch.

Here is a simplified version of the code — saved in a file unet.py :

I will not go into detail here, but rather just mention important design choices. It can be useful to view the architecture in repeating blocks in the encoder but also in the decoder path. As you can see in unet.py the DownBlock and the UpBlock help to build the architecture. Both use smaller helper functions that return the correct layer, depending on what arguments are passed , e.g. if a 2D ( dim=2 ) or 3D ( dim=3 ) network is wanted. The number of blocks is defined by the depth of the network.

A DownBlock generally has the following scheme:

A UpBlock has the following layers:

For our Unet class we just need to combine these blocks and make sure that the correct layers from the encoder are concatenated to the decoder (skip pathways). These layers have to be cropped if their sizes do not match with the corresponding layers from the decoder. In such cases, the autocrop function is used. At the end we just need to think about the parameter initialization. By default, the weights are initialized with torch.nn.init.xavier_uniform_ and the biases are initialized with zeros using torch.nn.init.zeros_ .

Creating a U-Net model

Let’s create such a model and use it to make a prediction on some random input:

from unet import UNet model = UNet(in_channels=1,

out_channels=2,

n_blocks=4,

start_filters=32,

activation='relu',

normalization='batch',

conv_mode='same',

dim=2)



x = torch.randn(size=(1, 1, 512, 512), dtype=torch.float32)

with torch.no_grad():

out = model(x)



print(f'Out: {out.shape}')

This will give us:

Out: torch.Size([1, 2, 512, 512])

To check weather our model is correct, we can get the model’s summary with:

from torchsummary import summary

summary(model, (1, 512, 512))

which prints out a summary like this:

A small note on the input sizes and padding: To ensure correct semantic concatenations, it is advised to use input sizes that return even spatial dimensions in every block but the last in the encoder. For example: An input size of 120² gives intermediate output shapes of [60², 30², 15²] in the encoder path for a U-Net with depth=4 . A U-Net with depth=5 with the same input size is not recommended, as a maxpooling operation on odd spatial dimensions (e.g. on a 15² input) should be avoided. Please also note that in the original publication, valid instead of same padding was used.

Now that we have built our model, it is time to create a training loop.