Quantization-aware training (QAT)

Quantization-aware training (QAT) simulates integer quantization during training so the model learns to compensate for quantization effects. This typically produces higher accuracy than post-training quantization while still generating an efficient integer model for on-device deployment.

For complete working examples, see the Tutorials page.

import dg

# Enable quantization on a trained model
qmodel = dg.enable_quantization(model, (x_train[:2],))

# Fine-tune with a lower learning rate
optimizer = torch.optim.Adam(qmodel.parameters(), lr=1e-4)
for epoch in range(epochs):
    for x, y in loader:
        optimizer.zero_grad()
        loss = F.cross_entropy(qmodel(x), y)
        loss.backward()
        optimizer.step()

# Freeze quantization
quantized = dg.freeze_quantization(qmodel)
schema = dg.export(quantized)

dg.enable_quantization takes a trained model and example inputs, returning the quantized graph in training mode. Use a lower learning rate (typically 1/10 of float training) for a few epochs, then freeze and export.