biapy.models.resunet_se

This module implements the Residual U-Net with Squeeze-and-Excitation (SE) blocks, referred to as ResUNet_SE.

This architecture is a variant of the U-Net, enhanced with residual connections and channel-wise attention mechanisms to improve feature learning and representation for various image analysis tasks.

The ResUNet_SE is designed for both 2D and 3D image processing, supporting different normalization types, kernel sizes, and flexible upsampling strategies. It can be configured for multi-task learning with multiple output heads, including an optional contrastive learning projection.

Key components and functionalities include:

Classes:

  • ResUNet_SE: The main Residual U-Net model with Squeeze-and-Excitation blocks.

This module relies on building blocks defined in biapy.models.blocks, such as ResConvBlock, ResUpBlock, ConvBlock, ProjectionHead, and normalization helpers.

Reference: Squeeze and Excitation Networks.

Image representation:

../../_images/unet.png

Image created with PlotNeuralNet.

class biapy.models.resunet_se.ResUNet_SE(image_shape=(256, 256, 1), activation='ELU', feature_maps=[32, 64, 128, 256], drop_values=[0.1, 0.1, 0.1, 0.1], normalization='none', k_size=3, upsample_layer='convtranspose', z_down=[2, 2, 2, 2], yx_down=[2, 2, 2, 2], output_channels=[1], separated_decoders=False, output_channel_info=['F'], explicit_activations: bool = False, head_activations: List[str] = ['ce_sigmoid'], upsampling_factor=(), upsampling_position='pre', isotropy=False, larger_io=True, extra_conv=True, contrast: bool = False, contrast_proj_dim: int = 256, return_one_tensor: bool = False)[source]

Bases: Module

Create Residual 2D/3D U-Net with Squeeze-and-Excitation (SE) blocks. Supports multi-head outputs and optional contrastive learning head.

forward(x) → Dict | Tensor[source]

Forward pass of the model.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, channels, height, width) for 2D or (batch_size, channels, depth, height, width) for 3D.

Returns:

Model output. Returns a dictionary if multi-head or contrastive outputs are enabled, otherwise returns the main prediction tensor.

Return type:

Dict or torch.Tensor