SegmentationHead¶

class SegmentationHead(in_channels, out_channels, kernel_size=3, activation=None, upsampling=1)[source]¶

Segmentation head for UNet++ architecture.

This class defines the final segmentation layer for the UNet++ model. It applies a convolution followed by optional upsampling and activation to produce the segmentation output.

conv2d¶

Convolutional layer for feature transformation.

Type:

nn.Conv2d

upsampling_layer¶

Upsampling layer (bilinear interpolation or identity).

Type:

nn.Module

activation¶

Activation function applied after upsampling.

Type:

nn.Module

Example

>>> head = SegmentationHead(in_channels=64, out_channels=2)
>>> x = torch.randn(1, 64, 128, 128)
>>> output = head(x)
>>> output.shape
... torch.Size([1, 2, 128, 128])

Initialize the SegmentationHead module.

This method sets up the segmentation head by creating a convolutional layer, an optional upsampling layer, and an activation function. It is typically used as the final stage in UNet++ architectures for semantic segmentation.

Parameters:

  • in_channels (int) – Number of input channels to the segmentation head.

  • out_channels (int) – Number of output channels (usually equal to the number of classes).

  • kernel_size (int) – Size of the convolution kernel. Defaults to 3.

  • activation (nn.Module | None) – Activation function applied after convolution. Defaults to None.

  • upsampling (int) – Upsampling factor applied to the output. Defaults to 1.

Methods

Attributes

training