Building a U-Net: layers definitions
In this and the next exercise, you will construct a U-Net architecture from scratch. Here, you start with defining the __init__() method where you will define layers and blocks of layers for the model to use.
The encoder and transposed convolution have already been defined for you. What's left is the decoder's convolutional blocks. You need to pass the appropriate number of input and output channels to each of them, taking into account the skip connections.
The first block, dec1, will take as input the concatenation of upconv3 output with the enc3 output. The dec1 output, in turn, should be equal to enc3 output. Can you fill-in all the missing input and output sizes?
This exercise is part of the course
Deep Learning for Images with PyTorch
Exercise instructions
- Define the three convolutional blocks in the decoder by passing the appropriate number of input and output channels to each of them.
Hands-on interactive exercise
Have a go at this exercise by completing this sample code.
class UNet(nn.Module):
def __init__(self, in_channels, out_channels):
super(UNet, self).__init__()
self.enc1 = self.conv_block(in_channels, 64)
self.enc2 = self.conv_block(64, 128)
self.enc3 = self.conv_block(128, 256)
self.enc4 = self.conv_block(256, 512)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
# Define the decoder blocks
self.dec1 = self.conv_block(____, ____)
self.dec2 = self.conv_block(____, ____)
self.dec3 = self.conv_block(____, ____)
self.out = nn.Conv2d(64, out_channels, kernel_size=1)