unet module
- class unet.DoubleConvolution(in_channels, out_channels)[source]
Bases:
ModuleApplies double convolution to the input
Double convolution is a sequence of two 3x3 convolutional layers with batch normalization and ReLU activation function.
- double_conv
Sequence of two convolutional layers with batch normalization and ReLU activation function
- Type:
nn.Sequential
- class unet.UNet(in_channels=3, out_channels=1, features=[64, 128, 256, 512])[source]
Bases:
ModuleUNet model
UNet is a convolutional neural network used for image segmentation. It consists of an encoder and a decoder. The encoder downsamples the input image and the decoder upsamples the encoder’s output to the original size. The encoder and decoder are connected by so-called skip connections, which concatenate the output of the encoder to the input of the decoder at the same resolution. This helps the decoder to recover the spatial information lost during downsampling. The skip connections are concatenated channel-wise, which means that the number of channels is doubled after each concatenation. The UNet model has a contracting path (encoder) and an expansive path (decoder): the contracting path follows the typical architecture of a convolutional neural network, with a series of convolutional layers followed by a max-pooling layer. The expansive path consists of a series of up-convolutions, which increase the spatial resolution of the input, followed by a series of convolutional layers. The final layer of the UNet model is a 1x1 convolutional layer that maps each pixel to the desired number of classes.
- num_features
Number of features in the encoder
- Type:
int
- ups
List of up-convolutions in the decoder
- Type:
nn.ModuleList
- downs
List of double convolutions in the encoder
- Type:
nn.ModuleList
- pool
Max pooling layer
- Type:
nn.MaxPool2d
- bottleneck
Bottleneck layer
- Type:
- final_conv
Final convolutional layer
- Type:
nn.Conv2d
References
Ronneberger, O., Fischer, P., & Brox, T. (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation. arXiv preprint arXiv:1505.04597.
- forward(x)[source]
Perform the forward propagation
- Parameters:
x (torch.Tensor) – Input tensor to be processed by the UNet model of shape (batch_size, in_channels, height, width)
- Raises:
ValueError – If the input tensor sizes is less then 2^(len(features)+1), since the tensor sizes are halved in each downsampling step
- Returns:
result – Output tensor after applying the UNet model
- Return type:
torch.Tensor