PyTorch Quantization
Quantization reduces the numerical precision of a model's weights and activations — from 32-bit floats (float32) to 8-bit integers (int8) or even 4-bit. This shrinks model size by up to 4×, speeds up inference by 2–4×, and cuts memory bandwidth usage — making it possible to run models on mobile phones, microcontrollers, and edge devices that could never run the full-precision version.
Why Quantization Works
Original model: Each weight stored as float32 (32 bits = 4 bytes) 1 million weights → 4 MB After int8 quantization: Each weight stored as int8 (8 bits = 1 byte) 1 million weights → 1 MB ← 4× smaller Speed improvement: Integer arithmetic is faster than floating-point on most hardware Especially on mobile chips with dedicated int8 accelerators
The Quantization Idea
float32 value: 0.347 → mapped to int8 value: 44 → stored as: 44 (1 byte instead of 4 bytes) During inference: int8 arithmetic is performed result is de-quantized back to float for the final output Accuracy loss: Rounding introduces small errors Well-quantized models lose < 1% accuracy on most tasks
Three Quantization Approaches in PyTorch
1. Post-Training Static Quantization (PTQ Static): → Train model normally → calibrate on unlabeled data → quantize → Fastest to apply, good accuracy → Requires calibration dataset 2. Post-Training Dynamic Quantization (PTQ Dynamic): → Quantize weights only, activations stay float32 → No calibration needed → simplest approach → Best for LSTM and Linear-heavy models 3. Quantization-Aware Training (QAT): → Simulate quantization during training → Model learns to be robust to quantization errors → Best accuracy, requires retraining
Dynamic Quantization — Easiest
Dynamic quantization quantizes only the weights ahead of time. Activations are quantized on-the-fly at runtime. No calibration data required — the simplest path to a smaller model.
import torch
import torch.nn as nn
import torch.quantization
model = nn.Sequential(
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 10)
)
# Quantize Linear layers to int8
quantized_model = torch.quantization.quantize_dynamic(
model,
{nn.Linear}, # which layer types to quantize
dtype=torch.qint8 # target precision
)
print(quantized_model)
x = torch.rand(8, 128)
out = quantized_model(x)
print(out.shape) # torch.Size([8, 10])Comparing Model Size Before and After
import os
import torch
def model_size_mb(model, path='tmp_model.pt'):
torch.save(model.state_dict(), path)
size_mb = os.path.getsize(path) / 1e6
os.remove(path)
return size_mb
original_size = model_size_mb(model)
quantized_size = model_size_mb(quantized_model)
print(f"Original: {original_size:.2f} MB")
print(f"Quantized: {quantized_size:.2f} MB")
print(f"Reduction: {original_size/quantized_size:.1f}×")Static Quantization — Better Accuracy
Static quantization also quantizes activations. It requires a calibration step where you run representative data through the model so PyTorch can measure activation ranges. The result is a fully int8 model — faster than dynamic quantization at inference time.
import torch
import torch.nn as nn
import torch.quantization
class QuantizableNet(nn.Module):
def __init__(self):
super().__init__()
self.quant = torch.quantization.QuantStub() # entry point
self.fc1 = nn.Linear(128, 64)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(64, 10)
self.dequant = torch.quantization.DeQuantStub() # exit point
def forward(self, x):
x = self.quant(x) # convert float to quantized
x = self.relu(self.fc1(x))
x = self.fc2(x)
x = self.dequant(x) # convert back to float
return x
model = QuantizableNet()
model.eval()
# Step 1: Set quantization configuration
model.qconfig = torch.quantization.get_default_qconfig('fbgemm') # x86 CPU
# Step 2: Prepare for calibration
torch.quantization.prepare(model, inplace=True)
# Step 3: Calibrate with representative data
calibration_data = torch.rand(100, 128)
with torch.no_grad():
model(calibration_data)
# Step 4: Convert to quantized model
torch.quantization.convert(model, inplace=True)
print("Static quantized model ready")
out = model(torch.rand(4, 128))
print(out.shape) # torch.Size([4, 10])Static Quantization Flow Diagram
Original float32 model
│
▼ torch.quantization.prepare()
Model with observer hooks
│
▼ Run calibration data (observe activation ranges)
Calibrated model
│
▼ torch.quantization.convert()
Quantized int8 model ← 4× smaller, 2–4× faster inference
Quantization-Aware Training (QAT)
QAT inserts fake quantization operations during training — the model sees quantization noise and learns weights that are robust to it. This produces the highest accuracy quantized model but requires retraining.
model = QuantizableNet()
model.train()
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
# Prepare for QAT (inserts fake quant nodes)
torch.quantization.prepare_qat(model, inplace=True)
# Normal training loop ...
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# ... train for several epochs ...
# Convert to quantized after training
model.eval()
torch.quantization.convert(model, inplace=True)Choosing a Quantization Strategy
Situation Strategy
───────────────────────────────────────────────────────────────
Quick size reduction, no retraining Dynamic quantization
Best speed, limited calibration data Static quantization
Highest accuracy, willing to retrain QAT
Deploying on mobile (ARM) Static + mobile backend ('qnnpack')
Deploying on x86 CPU server Static + 'fbgemm' backend
Summary
Quantization reduces model weights and activations from float32 to int8, making models 4× smaller and 2–4× faster at inference. Dynamic quantization is the simplest approach — one line of code, no calibration needed, great for LSTM and linear models. Static quantization requires a calibration step but quantizes both weights and activations for maximum speed. Quantization-Aware Training produces the best accuracy by training with simulated quantization noise. Choose the backend based on your hardware: fbgemm for x86 CPUs, qnnpack for mobile ARM processors. Quantization is essential for deploying large models on resource-constrained devices.