fuse_modules(model: nn.Module, modules_to_fuse: Sequence[Sequence[str]], inplace: bool = False)Collapse Conv/BN/ReLU (or Linear/ReLU) runs into single fused modules.
fuse_modules is a pre-processing step you run before
lucid.quantization.prepare in the static-PTQ workflow. It rewrites
the model so each named run of adjacent layers becomes one module: the first
member is replaced by the fused module and the rest by
lucid.nn.Identity (so the model's original forward still routes
through them harmlessly). BatchNorm is folded into the preceding convolution's
weight via the eval-time fusion, and a trailing ReLU is absorbed into an
nn.intrinsic fused module.
Fusion matters for quantization on two counts. First, folding BN into the
conv weight removes a separately-quantized layer and its rounding error, and
it means there is only one weight tensor to quantize per block. Second,
absorbing the ReLU makes the downstream activation observer see the post-
ReLU range, so the calibrated grid is not wasted on the negative half that
the ReLU would clip anyway — a direct accuracy win. Put the model in eval
mode first: the BN fold uses the frozen running statistics, which is only
correct in eval.
Parameters
modelnn.Moduleeval mode first — BN folding is
eval-time (it consumes the frozen running mean / variance).modules_to_fusesequence of name groups[["conv1", "bn1"], ["layer1.0.conv1", "layer1.0.bn1", "layer1.0.relu"]].
Supported runs: [Conv, BN], [Conv, BN, ReLU], [Conv, ReLU],
and [Linear, ReLU]; any other pattern raises ValueError.inplacebool= FalseTrue fuse and return model itself; if False (default) fuse
a copy.deepcopy and leave the original untouched.Returns
nn.ModuleThe fused model, with each run's first module replaced by the fused
module and the remaining members turned into Identity.
Notes
- Order-preserving: replacements are spliced in through
_modulesrather thansetattr, so aSequential's execution order is unchanged (setattrwould move a re-added key to the end). - The trailing members become
Identityrather than being deleted, so any hand-writtenforwardthat still calls them keeps working. - This is the eval / static-PTQ fusion. For quantization-aware training
use
fuse_modules_qat, which keeps Conv+BN trainable and folds BN per forward instead of freezing it. - Non-destructive by default (deep-copy).
Examples
>>> import lucid.nn as nn
>>> import lucid.quantization as Q
>>> class Block(nn.Module):
... def __init__(self):
... super().__init__()
... self.conv = nn.Conv2d(3, 8, 3)
... self.bn = nn.BatchNorm2d(8)
... self.relu = nn.ReLU()
... def forward(self, x):
... return self.relu(self.bn(self.conv(x)))
>>> m = Block().eval()
>>> fused = Q.fuse_modules(m, [["conv", "bn", "relu"]])
>>> type(fused.conv).__name__ # Conv+BN+ReLU -> one fused module
'ConvReLU2d'
>>> type(fused.bn).__name__, type(fused.relu).__name__
('Identity', 'Identity')
An unsupported run (here [BN, ReLU] with no leading Conv/Linear) raises:
>>> Q.fuse_modules(m.eval(), [["bn", "relu"]])
ValueError: fuse_modules: unsupported fusion pattern ['BatchNorm2d', 'ReLU']See Also
- lucid.quantization.fuse_modules_qat—Trainable Conv/BN fusion for QAT.
- lucid.quantization.prepare—Run this before inserting observers.
lucid.nn.utils.fusion.fuse_conv_bn_eval—The eval-time BN-into-conv fold.