Creating and training a U-Net model with PyTorch for 2D & 3D semantic segmentation: Model building [2/4] | by Johannes Schmidt | Towards Data Science

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. 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.
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.
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. For merging, I concatenate along the channel dimension (see Concatenate). Instead of transposed convolutions we could also use upsampling layers (interpolation methods) that are followed by a 1x1 or 3x3 convolution block to reduce the channel dimension. Using interpolation generally gets rid of the checkerboard artifact. For 3D input consider using trilinear interpolation.
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_.
For details and the available parameter options, I encourage you to take a look at the code. Feel free to change the code to your needs or expand e.g. the number of activation functions.
Let’s create such a model and use it to make a prediction on some random input:
from unet import UNetmodel = 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 this package pytorch-summary:
from torchsummary import summarysummary = summary(model, (1, 512, 512))
which prints out a summary like this:
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 1202 gives intermediate output shapes of [602, 302, 152] 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 152 input) should be avoided.
To make our lives easier, we can numerically compute the maximum network depth for a given input dimension with a simple function:
shape = 1920def compute_max_depth(shape, max_depth=10, print_out=True):    shapes = []    shapes.append(shape)    for level in range(1, max_depth):        if shape % 2 ** level == 0 and shape / 2 ** level > 1:            shapes.append(shape / 2 ** level)            if print_out:                print(f'Level {level}: {shape / 2 ** level}')        else:            if print_out:                print(f'Max-level: {level - 1}')            break    return shapesout = compute_max_depth(shape, print_out=True, max_depth=10)
This will output
Level 1: 960.0Level 2: 480.0Level 3: 240.0Level 4: 120.0Level 5: 60.0Level 6: 30.0Level 7: 15.0Max-level: 7
which tells us that that we can design a U-Net as deep as this without having to worry about semantic mismatches. Conversely, we can also numerically determine the possible input shapes dimensions for a given depth:
low = 10high = 512depth = 8def compute_possible_shapes(low, high, depth):    possible_shapes = {}    for shape in range(low, high + 1):        shapes = compute_max_depth(shape,                                   max_depth=depth,                                   print_out=False)        if len(shapes) == depth:            possible_shapes[shape] = shapes    return possible_shapespossible_shapes = compute_possible_shapes(low, high, depth)
This will output
{256: [256, 128.0, 64.0, 32.0, 16.0, 8.0, 4.0, 2.0], 384: [384, 192.0, 96.0, 48.0, 24.0, 12.0, 6.0, 3.0], 512: [512, 256.0, 128.0, 64.0, 32.0, 16.0, 8.0, 4.0]}
which tells us that we can have 3 different input shapes with such a level 8 U-Net architecture. But I dare to say that such a network with this input size is probably not useful in practice.
In this part we created a configurable UNet model for the purpose of semantic segmentation. Now that we have built our model, it is time to create a training loop in the next chapter.