prepare_qat(model: nn.Module, qconfig: QConfig | QConfigMapping | None = None, inplace: bool = False)Insert fake-quant modules for quantization-aware training.
prepare_qat is the entry point of the QAT workflow — prepare_qat →
fine-tune → convert. Where static PTQ only observes a frozen
float model, QAT makes the quantization grid part of the forward pass: this
swaps every weighted layer for its lucid.nn.qat counterpart (which
fake-quantizes both its weight and its output) and attaches a fake-quant to
every QuantStub. During the subsequent fine-tuning the network feels
the int8 grid on each step — via the straight-through estimator of
fake_quantize — and learns weights that minimize the rounding error,
typically recovering most of the accuracy a low-bit model would otherwise
lose. Reach for it when plain static PTQ drops too much accuracy (aggressive
bit widths, quantization-sensitive architectures).
Unlike prepare, the returned model is left in train mode: BN
keeps updating its running stats and the fake-quant grids track the evolving
ranges. Fuse Conv/BN with fuse_modules_qat first so the folded
ConvBn* modules stay trainable and fold BN per forward under the STE.
After fine-tuning, convert maps the nn.qat layers to the same
int8 inference layers a static-PTQ convert produces.
Parameters
modelnn.ModuleQuantStub / DeQuantStub boundaries (and,
ideally, Conv/BN runs already fused for QAT).QConfig is promoted to a global
mapping, and None (default) falls back to
get_default_qat_qconfig_mapping (fused weight + activation
fake-quant).inplacebool= FalseTrue mutate and return model itself; if False (default)
prepare a copy.deepcopy and leave the original untouched.Returns
nn.ModuleThe QAT-ready model in train mode. Fine-tune it, then call
convert to obtain the int8 inference model.
Notes
- Returned in
trainmode (contrastprepare, which returnseval) — fake-quant and BN both need training-time behavior. - Weighted layers become trainable fake-quant modules; their float weights
are kept and updated by the optimizer (nothing is baked to int8 until
convert). - Fuse with
fuse_modules_qat(notfuse_modules) beforehand: the eval-time BN fold would freeze BN, which QAT needs to keep learning. - Non-destructive by default (deep-copy).
Examples
>>> import lucid, lucid.nn as nn
>>> import lucid.quantization as Q
>>> model = nn.Sequential(nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10))
>>> model.qconfig = Q.get_default_qat_qconfig()
>>> qat = Q.prepare_qat(model) # fake-quant inserted, train mode
>>> qat.training
True
>>> # ... fine-tune qat for a few epochs with real labels ...
>>> qmodel = Q.convert(qat.eval()) # bake into the int8 inference modelSee Also
- lucid.quantization.prepare—Static-PTQ observers (no fake-quant, eval mode).
- lucid.quantization.convert—Bake the fine-tuned QAT model into int8.
- lucid.quantization.fuse_modules_qat—Trainable Conv/BN fusion for QAT.
- lucid.quantization.fake_quantize—The STE primitive QAT trains through.