SegmentationHead¶

class SegmentationHead(in_channels, out_channels, kernel_size=3)[source]¶

Segmentation head for UNet architecture.

This class defines the final segmentation layer for the UNet model. It applies a convolution to produce the segmentation output.

conv2d¶

Convolutional layer for feature transformation to output classes.

Type:

nn.Conv2d

Example

>>> head = SegmentationHead(in_channels=16, out_channels=1)
>>> x = torch.randn(1, 16, 224, 224)
>>> output = head(x)
>>> output.shape
... torch.Size([1, 1, 224, 224])

Initialize the SegmentationHead module.

This method sets up the segmentation head by creating a convolutional layer. 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.

Methods

Attributes

training