class

LazyBatchNorm2d

extends_LazyBatchNormMixin

LazyBatchNorm2d(eps: float = 1e-05, momentum: float | None = 0.1, affine: bool = True, track_running_stats: bool = True, device: DeviceLike = None, dtype: DTypeLike = None)

source

BatchNorm2d with lazy num_features inference.

Identical to BatchNorm2d except that num_features is inferred from x.shape[1] on the first forward pass.

Parameters

epsfloat= 1e-05

Numerical stability constant. Default: 1e-5.

momentumfloat or None= 0.1

EMA factor for running statistics. None uses cumulative averaging. Default: 0.1.

affinebool= True

If True, lazily allocates learnable scale and shift. Default: True.

track_running_statsbool= True

If True, lazily allocates running statistics buffers. Default: True.

deviceDeviceLike= None

Device for lazily allocated tensors. Default: None.

dtypeDTypeLike= None

Data type for lazily allocated tensors. Default: None.

Notes

Input: $(N, C, H, W)$ — same as BatchNorm2d. C is inferred on the first forward pass.
Output: same shape as the input.

Examples

>>> import lucid
>>> import lucid.nn as nn
>>> bn = nn.LazyBatchNorm2d()
>>> x = lucid.randn(8, 128, 32, 32)
>>> out = bn(x)    # num_features=128 inferred on first call
>>> out.shape
(8, 128, 32, 32)
Suitable for dynamic architectures where channel count depends on
a configuration value not available at module construction:
>>> def make_norm(affine: bool = True) -> nn.LazyBatchNorm2d:
...     return nn.LazyBatchNorm2d(affine=affine)