fuse_conv_bn_eval

→Module

fuse_conv_bn_eval(conv: object, bn: object)

source

Fold a BatchNorm layer into the preceding Conv weights (inference-only).

Because BN at eval time is an affine map with frozen statistics, the composition BN(Conv(x)) is itself a single convolution — the BN can be analytically absorbed into the conv's weight and bias. The fused module computes exactly the same output but with one fewer kernel launch and one fewer allocation per call. Standard step in deployment / quantisation pipelines.

Parameters

convConv1d | Conv2d | Conv3d

Convolution module whose weight / bias will absorb the BN.

bnBatchNorm1d | BatchNorm2d | BatchNorm3d

BatchNorm whose rank matches conv. Must be in eval mode (or otherwise be using its frozen running_mean / running_var); calling on a training-mode graph silently yields wrong outputs.

Returns

Module

Deep copy of conv with fused parameters. Originals are not mutated.

Raises

TypeError

If conv or bn is not one of the supported types.

Notes

Let $\mu$ , $\sigma^2$ , $\epsilon$ be the BN running statistics and $\gamma$ , $\beta$ its affine parameters (treated as $1$ and $0$ if affine=False). The fused weight and bias are

\mathbf{W}_{\text{fused}} \;=\; \mathbf{W} \cdot \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}},

b_{\text{fused}} \;=\; \frac{\gamma\,(b - \mu)}{\sqrt{\sigma^2 + \epsilon}} + \beta,

where the scale broadcasts along the output-channel axis. When the original conv has no bias, $b$ is taken as $0$ and the fused module gains one.

Examples

>>> import lucid.nn as nn
>>> from lucid.nn.utils import fuse_conv_bn_eval
>>> conv = nn.Conv2d(3, 16, 3); bn = nn.BatchNorm2d(16)
>>> conv.eval(); bn.eval()
>>> fused = fuse_conv_bn_eval(conv, bn)