FakeQuantizeCallback trains a model through the rounding it will be subject to. On every forward the weights are rounded onto the grid of the width you asked for and, when act_bits is set, so is every rounded layer’s output. The loss is computed on the rounded model, and the gradient reaches an untouched floating-point copy which the optimizer keeps training. When the fit ends the rounding is baked in, and what is left is an ordinary model holding rounded weights.
The accuracy it measures is simulated: the model keeps its float32 tensors, it is not smaller, it is not quicker, and nothing on this page measures a latency.
Imagenette at 160 pixels and one resnet18, fine-tuned for one cycle. fresh() gives every section that same checkpoint.
import io, torch
from fastai.vision.all import *
from fasterai.core.precision import fake_quant_spec
from fasterai.core.schedule import lin
from fasterai.quantize.fake_quantizer import FakeQuantizer
from fasterai.quantize.fake_quantize_callback import FakeQuantizeCallback
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'torch {torch.__version__} on '
f'{torch.cuda.get_device_name(0) if device == "cuda" else "cpu"}')torch 2.9.1+cu128 on NVIDIA GeForce RTX 5090path = untar_data(URLs.IMAGENETTE_160)
def make_dls():
"Imagenette at 160 pixels; `Normalize` is named here rather than left to `vision_learner`"
return ImageDataLoaders.from_folder(path, valid='val', item_tfms=Resize(160), bs=32,
batch_tfms=Normalize.from_stats(*imagenet_stats))
set_seed(42, reproducible=True)
dls = make_dls()
N = len(dls.valid_ds)
# fixed calibration set
files = sorted(get_image_files(path/'train'))
calib_dl = dls.test_dl(files[::len(files)//160][:160], bs=32)
print(f'{len(dls.train_ds)} training images, {N} validation images, '
f'calibration set of {len(calib_dl)} batches')9469 training images, 3925 validation images, calibration set of 5 batchesset_seed(42, reproducible=True)
learn = vision_learner(dls, resnet18, metrics=accuracy)
with learn.no_bar(), learn.no_logging():
learn.fine_tune(1)
CHECKPOINT = {k: v.detach().clone() for k, v in learn.model.state_dict().items()}
def fresh():
"That checkpoint in a learner of its own, unfrozen, with its loaders seeded the same way"
set_seed(0, reproducible=True)
l = vision_learner(make_dls(), resnet18, metrics=accuracy)
l.model.load_state_dict(CHECKPOINT)
l.unfreeze()
return l
FLOAT = score(learn)
report('fine-tuned, floating point', FLOAT)fine-tuned, floating point 96.28% (3779/3925) Wilson 95% [95.64, 96.83]FakeQuantizeCallback(weight_bits=4, act_bits=8) rounds every Conv2d and Linear weight to 4 bits, one scale per output channel, and the output of each of those layers to 8 bits. The fit is an ordinary fit_one_cycle: the callback says what it is training through when the fit starts, and what it baked when the fit ends.
Training through W4A8, weights per_channel
Baked in W4A8: the model holds its rounded weights, and fake_quantizer.remove() gives the trained floating-point ones back
W4A8, trained through it 92.59% (3634/3925) Wilson 95% [91.72, 93.36]print_precision() names the width every layer carries, and the spec the model comes out with says trained=True — which a post-training spec does not.
Simulated Precision Report:
--------------------------------------------------------------------------------
Layer Type Weight Act Weight axis
--------------------------------------------------------------------------------
0.0 Conv2d 4 bits 8 bits per_channel
0.4.0.conv1 Conv2d 4 bits 8 bits per_channel
0.4.0.conv2 Conv2d 4 bits 8 bits per_channel
0.4.1.conv1 Conv2d 4 bits 8 bits per_channel
0.4.1.conv2 Conv2d 4 bits 8 bits per_channel
0.5.0.conv1 Conv2d 4 bits 8 bits per_channel
0.5.0.conv2 Conv2d 4 bits 8 bits per_channel
0.5.0.downsample.0 Conv2d 4 bits 8 bits per_channel
0.5.1.conv1 Conv2d 4 bits 8 bits per_channel
0.5.1.conv2 Conv2d 4 bits 8 bits per_channel
0.6.0.conv1 Conv2d 4 bits 8 bits per_channel
0.6.0.conv2 Conv2d 4 bits 8 bits per_channel
0.6.0.downsample.0 Conv2d 4 bits 8 bits per_channel
0.6.1.conv1 Conv2d 4 bits 8 bits per_channel
0.6.1.conv2 Conv2d 4 bits 8 bits per_channel
0.7.0.conv1 Conv2d 4 bits 8 bits per_channel
0.7.0.conv2 Conv2d 4 bits 8 bits per_channel
0.7.0.downsample.0 Conv2d 4 bits 8 bits per_channel
0.7.1.conv1 Conv2d 4 bits 8 bits per_channel
0.7.1.conv2 Conv2d 4 bits 8 bits per_channel
1.4 Linear 4 bits 8 bits per_channel
1.8 Linear 4 bits 8 bits per_channel
--------------------------------------------------------------------------------
Overall W4A8
Rounded in floating point: every tensor is still stored at its original width, and the model holds a floating-point copy of every weight it rounds until remove().
spec W4A8 qscheme per_channel trained TrueThe same widths applied to the same checkpoint after training rather than during it: FakeQuantizer observes the activation ranges on the fixed calibration set, then rounds.
W4A8, rounded after training 89.96% (3531/3925) Wilson 95% [88.98, 90.86]
difference against the fit above: +103 validation images of 3925In this single run the fit that trained through the rounding ends 103 validation images of 3925 above the one rounded afterwards — 92.59% (3634/3925) against 89.96% (3531/3925). The two do not cost the same: the fit carries three epochs of training and this row carries none.
weight_schedule and weight_widths step the width down during the fit instead of rounding at the target from the first batch. The rungs share the Schedule’s progress equally, so three rungs over three epochs spend about one epoch each. act_schedule and act_widths do the same for the activations.
learn = fresh()
cb = FakeQuantizeCallback(weight_bits=4, act_bits=8, weight_schedule=lin, weight_widths=(16, 8, 4))
with learn.no_bar(), learn.no_logging():
learn.fit_one_cycle(3, 1e-3, cbs=[cb, ShowWidths(cb)])
LADDER = score(learn)
report('W4A8 down a 16 -> 8 -> 4 ladder', LADDER)
print(f'difference against the fit that started at 4 bits: {LADDER[1] - QAT[1]:+d} '
f'validation images of {N}')
spec = fake_quant_spec(learn.model)
print(f'spec {spec.label} weight_widths {spec.weight_widths}')Training through W4A8, weights per_channel, stepping weight through [16, 8, 4] bits
end of epoch 0: {'weight': 16, 'act': 8}
end of epoch 1: {'weight': 8, 'act': 8}
end of epoch 2: {'weight': 4, 'act': 8}
Baked in W4A8: the model holds its rounded weights, and fake_quantizer.remove() gives the trained floating-point ones back
W4A8 down a 16 -> 8 -> 4 ladder 92.28% (3622/3925) Wilson 95% [91.40, 93.07]
difference against the fit that started at 4 bits: -12 validation images of 3925
spec W4A8 weight_widths (16, 8, 4)The widths step down one rung per epoch, and the spec records the ladder that was asked for next to the width the model ended at. In this single run, on one architecture and one seed, the ladder ended 12 validation images of 3925 below the fit that started at 4 bits — 92.28% (3622/3925) against 92.59% (3634/3925). Whether lowering the width during a fit is worth doing is not measured here: at equal epochs a ladder also spends part of its training at a wider width, so the two fits differ in more than the transition.
An ordinary model. bake() runs on its own when the fit ends, so nothing has to be called; the weights are the rounded ones, still float32, and torch.save takes the model as it is.
still an ordinary module: Sequential, weights in torch.float32| Call | What it does |
|---|---|
FakeQuantizeCallback(weight_bits=4, act_bits=8) |
trains through W4A8; None on either width leaves those tensors in floating point |
weight_schedule=lin, weight_widths=(16, 8, 4) |
steps the weight width down during the fit, the rungs sharing the schedule’s progress equally; act_schedule and act_widths do the same for the activations |
fake_quant_spec(model) |
the widths the model carries, trained=True, and the ladder that was asked for |
cb.bake() / cb.strip() |
bake() runs when the fit ends; after a fit that raised, call one of them by hand |
Measured on: Imagenette-160, 9469 training images and 3925 validation images at 160 pixels, batch size 32, on the torch build and device printed at the top of the page. One
fine_tune(1)underset_seed(42, reproducible=True), then onefit_one_cycle(3, 1e-3)per section on the unfrozen model underset_seed(0, reproducible=True), and onelearn.validate()per row on the same 3925 images. Single run, one architecture, one seed: no dispersion is quoted anywhere on this page, and every difference quoted above is one run against one run. Every accuracy is simulated in floating point.
Schedule a width ladder reads its progress from