spyrit.core.nnet.Unet.final_block

Unet.final_block(in_channels: int, mid_channel: int, out_channels: int, kernel_size=3) → Sequential

Defines the final block of the U-Net model.

The final block is composed of three convolutional layers. The first two are followed by a ReLU activation function and a batch normalization layer. The last convolutional layer is only composed of a convolution operation.

Args:

in_channels (int): Number of channels in the input tensor.

mid_channel (int): Number of channels in the intermediate tensor (i.e.) the tensor after the first and second convolutional layers.

out_channels (int): Number of channels in the output tensor.

kernel_size (int or tuple, optional): Size of the three convolution kernels. It is directly passed to torch.nn.Conv2d .Defaults to 3.

Returns:

torch.nn.Sequential: A sequential block with three convolutional layers and two ReLU activation functions / batch normalization layers.