PyTorchUNetModelsZoo Module

Collection of PyTorch UNet model variants and convenience constructors.

class HMB.PyTorchUNetModelsZoo.DoubleConv(inChannels, outChannels)[source]

Bases: Module

Double convolution block used throughout the U-Net encoder and decoder.

Short summary:

Two consecutive 3x3 convolution layers each followed by normalization and ReLU activation used as a basic encoder/decoder building block.

Parameters:

  • inChannels (int) – Number of input channels.

  • outChannels (int) – Number of output channels.

Variables:

net (torch.nn.Sequential) – The internal sequential conv->norm->ReLU layers.

Returns:

Output feature map of shape [B, outChannels, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.ConfigConv(inChannels, outChannels, norm='batch', dropout=0.0, residual=False)[source]

Bases: Module

Configurable double-convolution block with selectable normalization, dropout and residual option.

Short summary:

A DoubleConv-like block with optional normalization type, dropout, and an optional residual connection when input and output channels match.

Parameters:

  • inChannels (int) – Number of input channels.

  • outChannels (int) – Number of output channels.

  • norm (str) – Normalization type: “batch” | “instance” | “none”.

  • dropout (float) – Dropout probability applied after the block.

  • residual (bool) – Whether to add a residual skip when shapes permit.

Variables:

  • net (torch.nn.Sequential) – Core conv->norm->ReLU operations.

  • dropout (torch.nn.Module|None) – Dropout layer when requested.

  • residual (bool) – Flag indicating if residual connection will be used.

Returns:

Processed feature map with same spatial size as input.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.ResidualBlock(inChannels, outChannels)[source]

Bases: Module

Residual convolutional block.

Short summary:

Two conv->BN->ReLU layers with an identity skip connection. A 1x1 projection is applied to the identity when the number of channels differs.

Parameters:

  • inChannels (int) – Number of input channels.

  • outChannels (int) – Number of output channels.

Variables:

  • conv2 (conv1,) – Convolution layers.

  • bn2 (bn1,) – BatchNorm layers.

  • proj (torch.nn.Conv2d|None) – 1x1 conv used when channel projection is needed.

Returns:

Output tensor with applied residual addition.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.AttentionGate(F_g, F_l, F_int)[source]

Bases: Module

Attention gating module for skip connections.

Short summary:

Lightweight gating unit that computes attention coefficients for encoder skip features using a gating signal from the decoder, improving focus on relevant spatial locations.

Parameters:

  • F_g (int) – Channels of the gating signal from the decoder.

  • F_l (int) – Channels of the skip connection from the encoder.

  • F_int (int) – Intermediate channel width inside the gate.

Variables:

psi (W_g, W_x,) – Internal 1x1 conv + BN branches used to compute attention.

Returns:

Reweighted skip features of shape matching the input skip map.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(g, x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.DepthwiseSeparableConv(inChannels, outChannels)[source]

Bases: Module

Depthwise separable convolution block.

Short summary:

Implements a depthwise convolution followed by a pointwise convolution, used to reduce parameter counts and computation in lightweight models.

Parameters:

  • inChannels (int) – Number of input channels.

  • outChannels (int) – Number of output channels.

Variables:
Returns:

Activated output tensor with outChannels channels.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.SEBlock(channels, reduction=16)[source]

Bases: Module

Squeeze-and-Excitation (SE) block.

Short summary:

Performs global channel-wise pooling followed by a small bottleneck MLP that produces per-channel scaling weights applied to the input.

Parameters:

  • channels (int) – Number of input/output channels.

  • reduction (int) – Reduction ratio for the bottleneck. Default 16.

Variables:
Returns:

Recalibrated tensor of same shape as input.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.MultiResBlock(inChans, outChans, alpha=1.67)[source]

Bases: Module

Multi-resolution convolution block capturing features at multiple receptive fields.

Short summary:

Parallel convolutions with kernel sizes 3x3, 5x5, and 7x7 followed by channel weighting to capture multi-scale context within a single block while controlling parameter growth.

Parameters:

  • inChans (int) – Input channel count.

  • outChans (int) – Output channel count (total across all paths).

  • alpha (float) – Scaling factor for path channel allocation. Default 1.67.

Variables:

  • conv3x3 (nn.Sequential) – 3x3 convolution path.

  • conv5x5 (nn.Sequential) – 5x5 convolution path (decomposed to two 3x3).

  • conv7x7 (nn.Sequential) – 7x7 convolution path (decomposed to three 3x3).

  • convShortcut (nn.Conv2d) – 1x1 shortcut for residual connection.

  • batchNorm (nn.BatchNorm2d) – Final batch normalization.

Returns:

Multi-resolution feature map with outChans channels.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.DenseBlock(inChans, numLayers=4, growthRate=32, bnSize=4)[source]

Bases: Module

Dense connectivity block with bottleneck layers for parameter efficiency.

Short summary:

Stacks multiple bottleneck layers where each layer receives feature maps from all preceding layers as input, promoting feature reuse and gradient flow.

Parameters:

  • inChans (int) – Input channel count.

  • numLayers (int) – Number of bottleneck layers in block. Default 4.

  • growthRate (int) – Channels added per layer. Default 32.

  • bnSize (int) – Bottleneck expansion factor. Default 4.

Variables:

layers (nn.ModuleList) – Sequential bottleneck layers with dense connectivity.

Returns:

Concatenated output of all layers with inChans + numLayers*growthRate channels.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.RecurrentConvLayer(channels, t=2)[source]

Bases: Module

Recurrent convolutional layer with internal state feedback.

Short summary:

Applies convolution repeatedly for T timesteps where each step receives feedback from its previous output, enabling iterative refinement of spatial features.

Parameters:

  • channels (int) – Input/output channel count.

  • t (int) – Number of recurrent iterations. Default 2.

Variables:

conv (nn.Conv2d) – Shared convolution kernel applied at each timestep.

Returns:

Refined feature map after T recurrent steps.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.ASPP(inChans, outChans, dilations=(1, 6, 12, 18))[source]

Bases: Module

Atrous Spatial Pyramid Pooling for multi-scale context aggregation.

Short summary:

Parallel dilated convolutions at multiple rates plus image pooling to capture objects at different scales within a single feature map.

Parameters:

  • inChans (int) – Input channel count.

  • outChans (int) – Output channel count after fusion.

  • dilations (Tuple[int]) – Dilation rates for parallel branches. Default (1, 6, 12, 18).

Variables:

  • conv1x1 (nn.Sequential) – 1x1 convolution branch (rate=1).

  • conv3x3_1..3 (nn.Sequential) – Dilated 3x3 convolutions at specified rates.

  • imagePool (nn.Sequential) – Global average pooling branch.

  • project (nn.Sequential) – Final projection and dropout.

Returns:

Context-enriched feature map with outChans channels.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.PatchEmbedding(inChans, embedDim=256, patchSize=4)[source]

Bases: Module

2D image to patch embedding with optional positional encoding.

Short summary:

Converts input image into non-overlapping patches via convolutional projection and flattens them into a sequence for transformer processing.

Parameters:

  • inChans (int) – Input channel count.

  • embedDim (int) – Embedding dimension per patch. Default 256.

  • patchSize (int) – Patch size (height=width). Default 4.

Variables:

proj (nn.Conv2d) – Convolutional patch projection layer.

Returns:

Patch sequence [B, N, embedDim] and (patchH, patchW).

Return type:

Tuple[torch.Tensor, Tuple[int, int]]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.TransformerBlock(embedDim, numHeads=8, mlpRatio=4.0, dropout=0.1)[source]

Bases: Module

Standard transformer encoder block with pre-normalization.

Short summary:

Multi-head self-attention followed by position-wise feed-forward network with layer normalization and residual connections before each sub-layer.

Parameters:

  • embedDim (int) – Embedding dimension.

  • numHeads (int) – Number of attention heads. Default 8.

  • mlpRatio (float) – Hidden dimension ratio for MLP. Default 4.0.

  • dropout (float) – Dropout probability. Default 0.1.

Variables:

  • norm2 (norm1,) – Pre-normalization layers.

  • attn (nn.MultiheadAttention) – Multi-head self-attention module.

  • mlp (nn.Sequential) – Two-layer MLP with GELU activation.

  • drop (nn.Dropout) – Dropout layer after MLP.

Returns:

Transformed sequence of same shape as input.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.CBAM(channels, reduction=16)[source]

Bases: Module

Convolutional Block Attention Module (channel + spatial attention).

Short summary:

Sequential channel attention followed by spatial attention to adaptively refine feature maps along both channel and spatial dimensions.

Parameters:

  • channels (int) – Input/output channel count.

  • reduction (int) – Reduction ratio for channel attention. Default 16.

Variables:

  • channelAttn (ChannelAttn) – Channel attention submodule.

  • spatialAttn (SpatialAttn) – Spatial attention submodule.

Returns:

Attention-refined feature map with same shape as input.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.ChannelAttn(channels, reduction=16)[source]

Bases: Module

Channel attention branch of CBAM using squeeze-and-excitation.

Short summary:

Global average and max pooling followed by shared MLP to compute channel-wise attention weights that rescale feature channels.

Parameters:

  • channels (int) – Input channel count.

  • reduction (int) – Reduction ratio for bottleneck MLP. Default 16.

Variables:

mlp (nn.Sequential) – Shared MLP for both pooling streams.

Returns:

Channel-refined feature map.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.SpatialAttn(kernelSize=7)[source]

Bases: Module

Spatial attention branch of CBAM using channel aggregation.

Short summary:

Channel-wise average and max pooling followed by convolution to compute spatial attention map that highlights important regions in feature maps.

Parameters:

kernelSize (int) – Convolution kernel size for spatial attention. Default 7.

Variables:

conv (nn.Conv2d) – Convolution to generate spatial attention map.

Returns:

Spatially-refined feature map.

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.UNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True)[source]

Bases: Module

Standard 2D U-Net architecture.

Short summary:

A typical encoder-decoder U-Net with four downsampling stages, a bottleneck, and four symmetric upsampling stages. Supports learned transpose convolution upsampling or bilinear upsampling followed by 1x1 projection.

Parameters:

  • inputChannels (int) – Number of input image channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Filters in the first stage. Default 64.

  • useConvTranspose2d (bool) – Use ConvTranspose2d when True; otherwise use bilinear upsample.

Variables:

  • center (enc1..enc4,) – Encoder and bottleneck blocks.

  • dec1..dec4 (up1..up4,) – Upsampling and decoder blocks.

  • finalConv (torch.nn.Conv2d) – 1x1 conv to map to logits.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.DynamicUNet(inputChannels=3, numClasses=2, baseChannels=64, depth=4, upMode='transpose', norm='batch', dropout=0.0, residual=False)[source]

Bases: Module

Configurable U-Net implementation with variable depth.

Short summary:

Generalized U-Net that allows customizing depth, normalization type, dropout, residual usage and upsampling mode. Useful for experimenting with architectures without duplicating code.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base filter count. Default 64.

  • depth (int) – Number of downsampling stages. Default 4.

  • upMode (str) – “transpose” or “bilinear” upsampling. Default “transpose”.

  • norm (str) – Normalization type: “batch”|”instance”|”none”. Default “batch”.

  • dropout (float) – Dropout probability applied inside config blocks. Default 0.0.

  • residual (bool) – Whether to use residual connections inside blocks. Default False.

Variables:

  • encs (nn.ModuleList) – Encoder blocks.

  • pools (nn.ModuleList) – Pooling layers.

  • ups (nn.ModuleList) – Upsampling layers.

  • decs (nn.ModuleList) – Decoder blocks.

  • finalConv (torch.nn.Conv2d) – 1x1 conv to logits.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.AttentionUNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True)[source]

Bases: Module

Attention U-Net variant with gating on skip connections.

Short summary:

Applies attention gating to encoder skip features before merging them into the decoder, improving localization and suppressing irrelevant responses in the skip maps.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base filter count. Default 64.

  • useConvTranspose2d (bool) – Use ConvTranspose2d for upsampling when True.

Variables:

  • center (enc1..enc4,) – Encoder/bottleneck blocks.

  • att1..att4 (AttentionGate) – Attention gating modules for skips.

  • dec1..dec4 (up1..up4,) – Decoder modules.

  • finalConv (torch.nn.Conv2d) – 1x1 conv mapping to logits.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.MobileUNet(inputChannels=3, numClasses=2, baseChannels=32, useConvTranspose2d=True)[source]

Bases: Module

Lightweight U-Net using depthwise separable convolutions.

Short summary:

Efficient UNet variant that replaces standard convolutions with depthwise separable convolutions to reduce parameters and compute.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base filter count. Default 32.

  • useConvTranspose2d (bool) – Use ConvTranspose2d for upsampling when True.

Variables:

  • center (enc1..enc4,) – Encoder and bottleneck sequences.

  • dec1..dec4 (up1..up4,) – Decoder modules.

  • finalConv (torch.nn.Conv2d) – 1x1 conv for logits.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.ResidualUNet(inputChannels, numClasses, baseChannels, useConvTranspose2d=True)[source]

Bases: Module

Residual U-Net variant built from ResidualBlock components.

Short summary:

A U-Net where encoder and decoder stages use residual blocks to ease optimization and improve gradient flow for deeper models.

Parameters:

  • inputChannels (int) – Number of input channels.

  • numClasses (int) – Number of output classes.

  • baseChannels (int) – Base number of filters.

  • useConvTranspose2d (bool) – Use ConvTranspose2d when True.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.SEUNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True, seReduction=16)[source]

Bases: Module

U-Net with Squeeze-and-Excitation modules applied to encoder and center.

Short summary:

Inserts SE blocks after encoder and center blocks to recalibrate channel-wise responses, improving representational expressiveness.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base number of filters. Default 64.

  • useConvTranspose2d (bool) – If True use learned transpose upsampling.

  • seReduction (int) – Reduction ratio for SE bottleneck. Default 16.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.ResidualAttentionUNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True)[source]

Bases: Module

Residual U-Net combined with Attention Gates.

Short summary:

Uses residual blocks in the encoder/decoder with attention gating applied to skip connections to guide feature fusion.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base number of filters. Default 64.

  • useConvTranspose2d (bool) – If True use ConvTranspose2d for upsampling.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.MultiResUNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True)[source]

Bases: Module

MultiResUNet with MultiRes blocks per stage.

Short summary:

Employs multi-resolution convolution blocks that capture features at multiple receptive sizes within each stage for richer multi-scale context.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base filter count. Default 64.

  • useConvTranspose2d (bool) – If True use ConvTranspose2d for upsampling.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.DenseUNet(inputChannels=3, numClasses=2, baseChannels=32, useConvTranspose2d=True)[source]

Bases: Module

DenseUNet integrating compact DenseBlocks and transition convolutions.

Short summary:

Uses DenseBlocks with 1x1 transition convolutions to manage channel growth, combining strengths of DenseNet connectivity with U-Net decoding.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base channel width. Default 32.

  • useConvTranspose2d (bool) – If True use ConvTranspose2d for upsampling.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.R2UNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True, t=2)[source]

Bases: Module

Recurrent Residual U-Net using recurrent convolution layers.

Short summary:

Integrates RecurrentConvLayer modules inside encoder stages to capture iterative spatial context while preserving residual optimization benefits.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of segmentation classes. Default 2.

  • baseChannels (int) – Base channel count. Default 64.

  • useConvTranspose2d (bool) – If True use ConvTranspose2d for upsampling.

  • t (int) – Number of recurrent iterations in RCL. Default 2.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.ASPPUNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True)[source]

Bases: Module

U-Net with ASPP at the bottleneck for expanded receptive field.

Short summary:

Places an Atrous Spatial Pyramid Pooling module at the bottleneck to aggregate multi-scale context using parallel dilated convolutions.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of classes. Default 2.

  • baseChannels (int) – Base number of filters. Default 64.

  • useConvTranspose2d (bool) – If True use ConvTranspose2d for upsampling.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.TransUNet(inputChannels=3, numClasses=2, baseChannels=64, embedDim=256, numHeads=8, numEncoders=2, useConvTranspose2d=True)[source]

Bases: Module

Transformer-enhanced U-Net with a small transformer stack at bottleneck.

Short summary:

Converts bottleneck conv maps into token sequences using PatchEmbedding, processes them via a transformer encoder stack, and projects back to convolutional feature maps for decoding.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of classes. Default 2.

  • baseChannels (int) – Base filter count. Default 64.

  • embedDim (int) – Transformer embedding dimension. Default 256.

  • numHeads (int) – Number of attention heads. Default 8.

  • numEncoders (int) – Number of transformer encoder layers. Default 2.

  • useConvTranspose2d (bool) – If True use ConvTranspose2d for upsampling.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.CBAMUNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True)[source]

Bases: Module

U-Net variant applying CBAM attention to encoder skip features.

Short summary:

Applies CBAM (channel + spatial attention) to encoder skip features before merging them into decoder stages to improve salient feature focus.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base filter count. Default 64.

  • useConvTranspose2d (bool) – If True use ConvTranspose2d for upsampling.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.EfficientUNet(inputChannels=3, numClasses=2, baseChannels=32, useConvTranspose2d=True)[source]

Bases: Module

Efficient wrapper around MobileUNet.

Short summary:

Thin wrapper that exposes the same API but uses MobileUNet internals to provide a lightweight segmentation model.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of classes. Default 2.

  • baseChannels (int) – Base channels for underlying MobileUNet. Default 32.

  • useConvTranspose2d (bool) – If True use ConvTranspose2d for upsampling.

Returns:

Logits tensor of shape [B, numClasses, H, W].

Return type:

torch.Tensor

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.BoundaryAwareUNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True, boundaryWeight=0.5)[source]

Bases: Module

Boundary-aware U-Net with explicit boundary detection branch.

Short summary:

Integrates parallel boundary detection pathway with auxiliary loss supervision using Sobel filters to enhance segmentation boundary precision and reduce boundary blurring common in standard encoder-decoder architectures.

Parameters:

  • inputChannels (int) – Number of input image channels. Default 3.

  • numClasses (int) – Number of segmentation classes. Default 2.

  • baseChannels (int) – Base filter count. Default 64.

  • useConvTranspose2d (bool) – Use ConvTranspose2d for upsampling when True.

  • boundaryWeight (float) – Weighting factor for boundary supervision loss. Default 0.5.

Variables:

  • encoder (nn.ModuleList) – Standard U-Net encoder blocks.

  • pools (nn.ModuleList) – Max pooling layers.

  • center (DoubleConv) – Bottleneck block.

  • upsamples (nn.ModuleList) – Upsampling modules.

  • decoders (nn.ModuleList) – Decoder convolutional blocks.

  • boundaryEncoder (nn.ModuleList) – Parallel encoder for boundary features.

  • boundaryCenter (DoubleConv) – Boundary bottleneck block.

  • boundaryUpsamples (nn.ModuleList) – Boundary upsampling modules.

  • boundaryDecoders (nn.ModuleList) – Boundary decoder blocks.

  • boundaryHead (nn.Conv2d) – Boundary prediction head (binary edge map).

  • segmentationHead (nn.Conv2d) – Final segmentation logits head.

Returns:

Segmentation logits and boundary map.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class HMB.PyTorchUNetModelsZoo.VNet(inputChannels=1, numClasses=1, baseChannels=16, noOfLevels=5, useConvTranspose2d=True)[source]

Bases: Module

V-Net style residual encoder-decoder adapted to 2D images.

Short summary:

Simple VNet-like architecture built from the existing ResidualBlock primitive. The module is intentionally compact and follows the same construction patterns as other models in this file (encoder ModuleList, decoder ups and decs).

Parameters:

  • inputChannels (int) – Number of input channels. Default 1.

  • numClasses (int) – Number of output classes. Default 1.

  • baseChannels (int) – Filters in the first stage. Default 16.

  • noOfLevels (int) – Number of encoder levels. Default 5.

  • useConvTranspose2d (bool) – Use ConvTranspose2d for learned upsampling when True.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

HMB.PyTorchUNetModelsZoo.CreateUNet(inputChannels=3, numClasses=2, baseChannels=64, depth=4, upMode='transpose', norm='batch', dropout=0.0, residual=False, modelType='dynamic')[source]

Extended factory that supports a large set of UNet variants by modelType.

Short summary:

Returns an instantiated UNet variant selected by modelType. The factory exposes common constructor arguments used across variants for convenience.

Parameters:

  • inputChannels (int) – Number of input channels. Default 3.

  • numClasses (int) – Number of output classes. Default 2.

  • baseChannels (int) – Base channel count for architectures. Default 64.

  • depth (int) – Depth parameter for dynamic variants. Default 4.

  • upMode (str) – “transpose” or “bilinear” upsampling. Default “transpose”.

  • norm (str) – Normalization mode for configurable blocks. Default “batch”.

  • dropout (float) – Dropout probability for configurable blocks. Default 0.0.

  • residual (bool) – Whether to use residual blocks where supported. Default False.

  • modelType (str) – Case-insensitive model selection string, e.g. “dynamic”.

Returns:

Instantiated UNet variant ready for training or inference.

Return type:

nn.Module

HMB.PyTorchUNetModelsZoo.GetUNetModel(modelName, inputChannels, numClasses, baseChannels=64)[source]

Factory function to instantiate a UNet variant based on the provided model name.

Parameters:

  • modelName (str) – Name of the UNet variant to instantiate.

  • inputChannels (int) – Number of input channels for the model.

  • numClasses (int) – Number of output classes for the model.

  • baseChannels (int) – Base number of channels for the model. Default is 64.

Returns:

An instance of the requested UNet variant.

Return type:

nn.Module