spyrit.core.nnet.Unet.bottle_neck

Unet.bottle_neck(in_channels: int, kernel_size=3, padding=1) → Sequential

Defines the bottleneck block of the U-Net model.

The bottleneck block is composed of two convolutional layers, each followed by a ReLU activation function. The number of output channels of the first convolutional layer is twice the number of its input channels.

Args:

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

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

padding (int or tuple or string, optional): Input padding. It is directly passed to torch.nn.Conv2d, see its documentation for valid options. Defaults to 1.

Returns:

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