Implementation:NVIDIA DALI EfficientNet Model

Knowledge Sources	NVIDIA_DALI
Domains	Vision, Training
Last Updated	2026-02-08 16:00 GMT

Overview

Implements the full EfficientNet family of image classification models (B0-B7) in PyTorch with support for width/depth scaling, quantization, and TensorRT optimization.

Description

This module provides the complete EfficientNet architecture following the compound scaling methodology from the original paper. The EffNetArch dataclass defines the architecture specification including block type, stem channels, feature channels, per-stage kernel sizes, strides, repeat counts, expansion ratios, and output channels. The scale method applies compound width and depth scaling coefficients to produce the B0 through B7 variants from a single base configuration.

The EfficientNet class (inheriting from nn.Module) constructs the network from a stem (3x3 conv + BN + activation), a stack of Mobile Inverted Bottleneck (MBConv) blocks organized into stages, a feature head (1x1 conv + BN + activation), and a classifier head (adaptive average pooling + dropout + fully connected). The MBConvBlock implements the inverted bottleneck with optional expansion, depthwise separable convolution, Squeeze-and-Excitation attention, projection, and stochastic depth residual connections.

Two MBConv variants are provided: original_mbconv which computes SE squeeze dimensions from input channels, and widese_mbconv which computes them from the expanded hidden dimensions (the "Wide SE" variant used in NVIDIA's pretrained models). The module ships with pre-defined configurations for all 24 model variants (B0-B7 in original, Wide SE, and quantized forms), complete with NGC pretrained checkpoint URLs. All configurations are registered in the architectures dictionary for convenient access.

Usage

Use this module as the model backbone in the EfficientNet PyTorch training example with DALI data loading. Select a model variant from the architectures dictionary or use the EffNetParams dataclass to configure custom parameters for dropout, quantization, batch norm settings, and stochastic depth survival probability.

Code Reference

Source Location

• Repository: NVIDIA_DALI
• File: docs/examples/use_cases/pytorch/efficientnet/image_classification/models/efficientnet.py
• Lines: 1-577

Signature

@dataclass
class EffNetArch(ModelArch):
    block: Any
    stem_channels: int
    feature_channels: int
    kernel: List[int]
    stride: List[int]
    num_repeat: List[int]
    expansion: List[int]
    channels: List[int]
    default_image_size: int
    squeeze_excitation_ratio: float = 0.25
    def scale(self, wc, dc, dis, divisor=8) -> "EffNetArch": ...

@dataclass
class EffNetParams(ModelParams):
    dropout: float
    num_classes: int = 1000
    activation: str = "silu"
    conv_init: str = "fan_in"
    bn_momentum: float = 0.01
    bn_epsilon: float = 1e-3
    survival_prob: float = 1
    quantized: bool = False
    trt: bool = False

class EfficientNet(nn.Module):
    def __init__(self, arch, dropout, num_classes=1000, activation="silu",
                 conv_init="fan_in", bn_momentum=0.01, bn_epsilon=1e-3,
                 survival_prob=1, quantized=False, trt=False): ...
    def forward(self, x): ...
    def extract_features(self, x, layers=None): ...

class MBConvBlock(nn.Module):
    def __init__(self, builder, depsep_kernel_size, in_channels, out_channels,
                 expand_ratio, stride, squeeze_excitation_ratio,
                 squeeze_hidden=False, survival_prob=1.0,
                 quantized=False, trt=False): ...
    def forward(self, x): ...

Import

from image_classification.models.efficientnet import (
    EfficientNet, EffNetArch, EffNetParams, architectures
)

I/O Contract

Inputs (EfficientNet.__init__)

Name	Type	Required	Description
arch	EffNetArch	Yes	Architecture specification defining block types, channels, and scaling.
dropout	float	Yes	Dropout probability before the final FC classifier layer.
num_classes	int	No	Number of output classification classes. Default: 1000.
activation	str	No	Activation function: "silu", "relu", or "onnx-silu". Default: "silu".
conv_init	str	No	Kaiming initialization mode: "fan_in" or "fan_out". Default: "fan_in".
bn_momentum	float	No	Batch normalization momentum. Default: 0.01.
bn_epsilon	float	No	Batch normalization epsilon. Default: 1e-3.
survival_prob	float	No	Stochastic depth survival probability (1.0 = no dropping). Default: 1.
quantized	bool	No	Enable quantization-aware training. Default: False.
trt	bool	No	Use TensorRT-compatible SE blocks. Default: False.

Inputs (forward)

Name	Type	Required	Description
x	torch.Tensor	Yes	Input image tensor of shape [N, 3, H, W]. Image size depends on the variant (224 for B0, up to 600 for B7).

Outputs (forward)

Name	Type	Description
logits	torch.Tensor	Classification logits of shape [N, num_classes].

Available Model Variants

Variant	Width Coeff	Depth Coeff	Image Size	Dropout
efficientnet-b0	1.0	1.0	224	0.2
efficientnet-b1	1.0	1.1	240	0.2
efficientnet-b2	1.1	1.2	260	0.3
efficientnet-b3	1.2	1.4	300	0.3
efficientnet-b4	1.4	1.8	380	0.4
efficientnet-b5	1.6	2.2	456	0.4
efficientnet-b6	1.8	2.6	528	0.5
efficientnet-b7	2.0	3.1	600	0.5

Also available with efficientnet-widese-* (Wide SE) and efficientnet-quant-* (quantized) prefixes.

Usage Examples

Instantiating EfficientNet-B0

from image_classification.models.efficientnet import architectures

# Get model configuration and build
model_config = architectures["efficientnet-b0"]
model = model_config.constructor(
    arch=model_config.arch,
    **vars(model_config.params)
)

# Forward pass
import torch
x = torch.randn(8, 3, 224, 224)
logits = model(x)
print(logits.shape)  # torch.Size([8, 1000])

Extracting intermediate features

# Extract features from specific layers
features = model.extract_features(x, layers=["layer3", "layer5", "features"])
for name, feat in features.items():
    print(f"{name}: {feat.shape}")

Related Pages

• Environment:NVIDIA_DALI_CUDA_GPU_Environment