PythiloLesson 29 of 30 · progress kept in this browser

Quantization

A float32 weight takes four bytes. Int8 quantization stores each as one byte plus a scale shared by the whole tensor: q = round(x / scale) and back, approximately, x = q * scale. With the symmetric choice scale = max(abs(x)) / 127 the largest weight maps to 127 and the error of any weight is at most half a scale.

Dynamic quantization stores a model's linear weights in int8 and quantizes each batch of activations on the fly. On iPhone and iPad the QNNPACK engine runs the int8 arithmetic.

import io
import torch
from torch import nn
import torch.ao.quantization as tq

torch.backends.quantized.engine = "qnnpack"
model = nn.Sequential(nn.Linear(256, 256),
                      nn.ReLU(),
                      nn.Linear(256, 10))
small = tq.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8)

def size(m):
    buffer = io.BytesIO()
    torch.save(m.state_dict(), buffer)
    return buffer.tell()

print(size(model), size(small))
print(small[0])

The model gets about four times smaller and its outputs move only slightly: weights trained in float32 rarely need more than 256 levels.

Your task

Write quantize(x) returning (q, scale): symmetric int8 quantization with scale = max(abs(x)) / 127 (use 1.0 when x is all zeros) and q an int8 tensor. Write dequantize(q, scale) returning floats. Write quantize_model(model) returning a dynamically quantized copy of model, its nn.Linear layers in torch.qint8, using the qnnpack engine.

This lesson uses torch, which Python in this browser does not include. Read it here, and run and check it in the Pythilo app on iPhone or iPad.