EfficentNet3D Encoder

This encoder is currently used by EffUNet.

This file contains helper functions for building the model and for loading model parameters. These helper functions are built to mirror those in the official TensorFlow implementation.

3D EfficientNet adapted from: https://github.com/shijianjian/EfficientNet-PyTorch-3D

Usage example:

model = EfficientNet3D.from_name("efficientnet-b1", override_params={'include_top': False}, in_channels=1)
model.cuda()  # On CUDA machine
model.to('mps')  # On Apple Silicon

List of pyramid layers:

{
    0: ['_conv_stem'],              # 100
    1: ['_blocks', '1', '_bn0'],    # 50
    2: ['_blocks', '3', '_bn0'],    # 25
    3: ['_blocks', '5', '_bn0'],    # 12
    4: ['_blocks', '11', '_bn0'],   # 6
    5: ['_bn1']                     # 3
}

class biom3d.models.encoder_efficientnet3d.BlockDecoder

Block Decoder for EfficientNet blocks. From Tensorflow repository.

Provides utilities to encode/decode block configuration strings into BlockArgs namedtuples for better readability and processing.

static decode(string_list: list[str]) → list[BlockArgs]

Decode a list of block strings into a list of BlockArgs.

Parameters:

string_list (list of str) – List of block description strings.

Returns:

List of decoded block arguments.

Return type:

list of BlockArgs

static encode(blocks_args: list[BlockArgs]) → list[str]

Encode a list of BlockArgs into a list of strings.

Parameters:

blocks_args (list of BlockArgs) – List of block arguments to encode.

Returns:

List of encoded block description strings.

Return type:

list of str

class biom3d.models.encoder_efficientnet3d.EfficientNet3D(*args: Any, **kwargs: Any)

An EfficientNet model. Most easily loaded with the .from_name or .from_pretrained methods.

Variables:

• _blocks (nn.ModuleList) – List of MBConvBlock3D blocks making up the network.

• _conv_stem (nn.Conv3d) – Initial convolutional stem layer.

• _conv_head (nn.Conv3d) – Final convolutional head layer before classifier.

• _bn0 (nn.InstanceNorm3d) – Batch normalization after the stem.

• _bn1 (nn.BatchNorm3d) – Batch normalization after the head.

• _avg_pooling (nn.AdaptiveAvgPool3d) – Global average pooling layer.

• _dropout (nn.Dropout) – Dropout layer before the classifier.

• _fc (nn.Linear) – Fully connected linear layer for classification.

• _swish (nn.Module) – Swish activation function module.

Examples

>>> model = EfficientNet3D.from_name('efficientnet-b0')
>>> x = torch.randn(1, 3, 224, 224, 224)
>>> logits = model(x)

__init__(blocks_args: BlockArgs, global_params: GlobalParams | None = None, in_channels: int = 3, num_pools: list[int] = [5, 5, 5], first_stride: list[int] = [1, 1, 1])

Efficientnet 3D model implementation.

Parameters:

• blocks_args (list of BlockArgs) – List of block argument namedtuples to construct the network blocks.

• global_params (namedtuple) – Global parameters shared across blocks (e.g., dropout rate, batch norm params).

• in_channels (int, optional) – Number of input channels. Default is 3.

• num_pools (list of int, optional) – List defining the number of pooling layers per dimension. Default is [5, 5, 5].

• first_stride (list of int, optional) – Stride size of the first convolution layer in each spatial dimension. Default is [1, 1, 1].

Raises:

AssertionError – If block_args is not a list or empty

extract_features(inputs: torch.Tensor) → torch.Tensor

Extract features from inputs by forwarding through convolutional layers.

Parameters:

inputs (torch.Tensor) – Input tensor of shape (batch_size, channels, depth, height, width).

Returns:

Feature tensor after convolutional blocks and head.

Return type:

torch.Tensor

forward(inputs: torch.Tensor) → torch.Tensor

Forward pass of the EfficientNet3D model. Apply final Linear layer.

Parameters:

inputs (torch.Tensor) – Input tensor of shape (batch_size, channels, depth, height, width).

Returns:

Logits tensor of shape (batch_size, num_classes) if include_top=True, otherwise feature tensor.

Return type:

torch.Tensor

classmethod from_name(model_name: str, override_params: dict | None = None, in_channels: int = 3, **kwargs) → EfficientNet3D

Create an EfficientNet3D model from a predefined model name.

Parameters:

• model_name (str) – Name of the EfficientNet model variant, e.g. ‘efficientnet-b0’.

• override_params (dict or None, optional) – Dictionary to override global parameters.

• in_channels (int, optional) – Number of input channels.

• **kwargs (dict) – Additional keyword arguments passed to the constructor.

Returns:

Constructed EfficientNet3D model.

Return type:

EfficientNet3D

classmethod get_image_size(model_name: str) → int

Get the default input image size for a given EfficientNet model.

Parameters:

model_name (str) – EfficientNet model name.

Returns:

Default image size.

Return type:

int

set_swish(memory_efficient: bool = True) → None

Set the Swish activation function implementation.

Parameters:

memory_efficient (bool, optional) – If True, use memory-efficient Swish (suitable for training). If False, use standard Swish (suitable for export). Default is True.

Return type:

None

class biom3d.models.encoder_efficientnet3d.MBConvBlock3D(*args: Any, **kwargs: Any)

Mobile Inverted Residual Bottleneck Block.

Variables:

has_se (bool) – Whether the block contains a Squeeze and Excitation layer.

__init__(block_args: BlockArgs, global_params: GlobalParams)

Mobile Inverted Residual Bottleneck Block (3D).

Parameters:

• block_args (namedtuple) – BlockArgs namedtuple containing block configuration.

• global_params (namedtuple) – GlobalParams namedtuple containing global model parameters.

forward(inputs: torch.Tensor, drop_connect_rate: float | None = None) → torch.Tensor

Forward pass through the MBConv block.

Parameters:

• inputs (torch.Tensor) – Input tensor to the block.

• drop_connect_rate (float, optional) – Drop connect probability (between 0 and 1). Default is None.

Returns:

Output tensor after processing through the block.

Return type:

torch.Tensor

set_swish(memory_efficient: bool = True) → None

Set the Swish activation function implementation.

Parameters:

memory_efficient (bool, optional) – If True, use memory-efficient Swish (suitable for training). If False, use standard Swish (suitable for export). Default is True.

Return type:

None

class biom3d.models.encoder_efficientnet3d.MemoryEfficientSwish(*args: Any, **kwargs: Any)

Module wrapping the memory-efficient Swish activation function.

forward(x: torch.Tensor) → torch.Tensor

Apply the memory-efficient Swish activation.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor after Swish activation.

Return type:

torch.Tensor

class biom3d.models.encoder_efficientnet3d.Swish(*args: Any, **kwargs: Any)

Standard Swish activation function module.

forward(x: torch.Tensor) → torch.Tensor

Apply Swish activation function: x * sigmoid(x).

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor after Swish activation.

Return type:

torch.Tensor

class biom3d.models.encoder_efficientnet3d.SwishImplementation(*args: Any, **kwargs: Any)

Memory-efficient implementation of the Swish activation function.

Swish function: x * sigmoid(x)

static backward(ctx, grad_output: torch.Tensor) → torch.Tensor

Backward pass for Swish activation.

Parameters:

• ctx (Context object) – Contains saved tensors from forward pass.

• grad_output (torch.Tensor) – Gradient of the loss with respect to the output.

Returns:

Gradient of the loss with respect to the input.

Return type:

torch.Tensor

static forward(ctx, i: torch.Tensor) → torch.Tensor

Forward pass for Swish activation.

Parameters:

• ctx (Context object) – Used to stash information for backward computation.

• i (torch.Tensor) – Input tensor.

Returns:

Output tensor after applying Swish.

Return type:

torch.Tensor

biom3d.models.encoder_efficientnet3d.drop_connect(inputs: torch.Tensor, p: float, training: bool) → torch.Tensor

Apply drop connect (stochastic depth) regularization.

Parameters:

• inputs (torch.Tensor) – Input tensor of shape (batch_size, …).

• p (float) – Drop connect probability.

• training (bool) – Whether the model is in training mode.

Returns:

Tensor after applying drop connect.

Return type:

torch.Tensor

biom3d.models.encoder_efficientnet3d.efficientnet3d(width_coefficient: float | None = None, depth_coefficient: float | None = None, dropout_rate: float = 0.2, drop_connect_rate: float = 0.2, image_size: int | None = None, num_classes: int = 1000, include_top: bool = True) → tuple[BlockArgs, GlobalParams]

Create EfficientNet3D block arguments and global parameters.

Parameters:

• width_coefficient (float, optional) – Width multiplier to scale number of filters.

• depth_coefficient (float, optional) – Depth multiplier to scale number of layers.

• dropout_rate (float, default=0.2) – Dropout rate before the classifier.

• drop_connect_rate (float, default=0.2) – Drop connect rate for stochastic depth.

• image_size (int, optional) – Input image size (default is None).

• num_classes (int, default=1000) – Number of output classes (default is 1000).

• include_top (bool, default=True) – Whether to include the classification head (default is True).

Returns:

• blocks_args (list of BlockArgs) – List of block arguments describing the network architecture.

• global_params (GlobalParams) – Namedtuple of global model parameters.

biom3d.models.encoder_efficientnet3d.efficientnet_params(model_name: str) → tuple[float, float, int, float]

Map EfficientNet model name to parameter coefficients.

Parameters:

model_name (str) – Name of the EfficientNet model.

Returns:

A tuple with width_coefficient, depth_coefficient, image_size, dropout_rate.

Return type:

tuple of (float, float, int, float)

biom3d.models.encoder_efficientnet3d.get_model_params(model_name: str, override_params: dict | None) → tuple[BlockArgs, GlobalParams]

Retrieve EfficientNet block args and global parameters by model name.

Parameters:

• model_name (str) – EfficientNet model name (e.g., ‘efficientnet-b0’).

• override_params (dict or None) – Dictionary to override default global parameters.

Raises:

• NotImplementedError – If model name doesn’t start with ‘efficientnet’

• ValueError – If override_params has field not in GlobalParams

Returns:

• blocks_args (list of BlockArgs) – List of block arguments.

• global_params (GlobalParams) – Namedtuple of global parameters.

biom3d.models.encoder_efficientnet3d.round_filters(filters: int, global_params: GlobalParams) → int

Calculate and round number of filters based on depth multiplier.

Parameters:

• filters (int) – Original number of filters.

• global_params (GlobalParams) – Global parameters containing width_coefficient, depth_divisor, and min_depth.

Returns:

Rounded (floor) number of filters adjusted by width coefficient and divisor.

Return type:

int

biom3d.models.encoder_efficientnet3d.round_repeats(repeats: int, global_params: GlobalParams) → int

Round the number of block repeats ased on depth multiplier..

Parameters:

• repeats (int) – Original number of repeats.

• global_params (GlobalParams) – Global parameters containing depth_coefficient.

Returns:

Rounded (ceil) number of repeats adjusted by depth coefficient.

Return type:

int

biom3d.models.encoder_efficientnet3d.GlobalParams(batch_norm_momentum, batch_norm_epsilon, dropout_rate, num_classes, width_coefficient, depth_coefficient, depth_divisor, min_depth, drop_connect_rate, image_size, include_top)

Global model parameters used across the entire network (stem, blocks, and head).

Parameters:

• batch_norm_momentum (float) – Momentum for batch normalization layers.

• batch_norm_epsilon (float) – Small epsilon value for numerical stability in batch norm.

• dropout_rate (float) – Dropout rate used before the classifier head.

• num_classes (int) – Number of output classes for classification.

• width_coefficient (float) – Width scaling coefficient for number of filters.

• depth_coefficient (float) – Depth scaling coefficient for number of block repeats.

• depth_divisor (int) – Divisor to ensure the number of filters is divisible by this value.

• min_depth (int) – Minimum depth (number of filters).

• drop_connect_rate (float) – Drop connect probability for stochastic depth.

• image_size (int) – Input image size.

• include_top (bool) – Whether to include the fully-connected classifier head.

biom3d.models.encoder_efficientnet3d.BlockArgs(kernel_size, num_repeat, input_filters, output_filters, expand_ratio, id_skip, stride, se_ratio)

Named tuple describing the parameters for a single model block.

Parameters:

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

• num_repeat (int) – Number of block repeats.

• input_filters (int) – Number of input filters.

• output_filters (int) – Number of output filters.

• expand_ratio (int) – Expansion ratio for internal layers.

• id_skip (bool) – Whether to use skip connections.

• stride (int) – Stride size.

• se_ratio (float or None) – Squeeze-and-excitation ratio.