fuse_modules_qat(model: nn.Module, modules_to_fuse: Sequence[Sequence[str]], qconfig: QConfig, inplace: bool = False)Fuse for QAT — Conv+BN runs become trainable nn.intrinsic.qat modules.
fuse_modules_qat is the fusion step of the quantization-aware-training
workflow, run before lucid.quantization.prepare_qat. It mirrors
fuse_modules but keeps BatchNorm learnable. Where the eval-time
fuse_modules folds BN into the conv weight once and freezes it — fine
for a frozen PTQ model — that would stop BN from training, which QAT needs.
So here a [Conv, BN] / [Conv, BN, ReLU] run becomes a trainable
ConvBn* / ConvBnReLU* module that folds BN per forward: the BN
parameters keep updating and the fold is recomputed each step, all under the
straight-through estimator of the fake-quant grid. [Conv, ReLU] and
[Linear, ReLU] runs (no BN) become the plain nn.intrinsic fused
module, which lucid.quantization.prepare_qat then swaps to its QAT
form.
Because the trainable ConvBn* modules fake-quantize at construction time,
the qconfig must be supplied here rather than later — it is baked into
the fused module so its weight fake-quant grid is active from the first
forward. As with fuse_modules, the first member of each run is
replaced by the fused module and the rest become Identity.
Parameters
modelnn.Modulefuse_modules, do
not put it in eval — BN must keep learning).modules_to_fusesequence of name groupsfuse_modules — e.g.
[["conv1", "bn1"], ["layer1.0.conv1", "layer1.0.bn1"]]. Same supported
runs: [Conv, BN], [Conv, BN, ReLU], [Conv, ReLU],
[Linear, ReLU].qconfigQConfigConvBn* modules; they fake-quant
at construction, so it is needed at fusion time rather than at
lucid.quantization.prepare_qat time.inplacebool= FalseTrue fuse and return model itself; if False (default) fuse
a copy.deepcopy and leave the original untouched.Returns
nn.ModuleThe fused, still-trainable model — hand it to
lucid.quantization.prepare_qat, fine-tune, then
lucid.quantization.convert.
Notes
- Conv+BN stays a single trainable module that folds BN per forward; BN
is never frozen (contrast
fuse_modules, which folds once at eval). - The
qconfigargument is mandatory precisely because the fusedConvBn*fake-quantizes at construction — there is no later hook to pass it through. - Order-preserving splice via
_modules; trailing members becomeIdentityso the originalforwardstill routes through them. - 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)
... def forward(self, x):
... return self.bn(self.conv(x))
>>> m = Block() # kept in train mode
>>> qcfg = Q.get_default_qat_qconfig()
>>> fused = Q.fuse_modules_qat(m, [["conv", "bn"]], qcfg)
>>> type(fused.conv).__name__ # trainable Conv+BN, folds per forward
'ConvBn2d'
>>> qat = Q.prepare_qat(fused) # continue the QAT workflowSee Also
- lucid.quantization.fuse_modules—Eval-time (static-PTQ) fusion; freezes BN.
- lucid.quantization.prepare_qat—Run this after fusing, before fine-tuning.
- lucid.quantization.convert—Bake the fine-tuned QAT model into int8.