qat.py

import argparse
from datetime import datetime

import torch
import torch.nn as nn
from torchao.quantization import quantize_
from torchao.quantization.granularity import PerAxis, PerToken
from torchao.quantization.qat import IntxFakeQuantizeConfig, QATConfig, QATStep
from torchao.quantization.quant_primitives import MappingType

from model import Device, MNISTModel
from train import _save_model, train


def qat(
    model: nn.Module,
    output: str,
    epochs: int,
    lr: float,
    batch: int,
    device: Device = "cpu",
) -> None:
    model = model.to(device)

    # prepare: swap `torch.nn.Linear` -> `FakeQuantizedLinear`
    activation_config = IntxFakeQuantizeConfig(
        dtype=torch.uint8,
        granularity=PerToken(),
        mapping_type=MappingType.ASYMMETRIC,
    )

    weight_config = IntxFakeQuantizeConfig(
        dtype=torch.int8,
        granularity=PerAxis(0),
        mapping_type=MappingType.SYMMETRIC,
    )

    qat_config_prepare = QATConfig(
        activation_config=activation_config,
        weight_config=weight_config,
        step=QATStep.PREPARE,
    )

    # prepare: swap `torch.nn.Linear` -> `FakeQuantizedLinear`
    quantize_(model, qat_config_prepare)

    # train: train the model with QAT (don't save weights yet)
    train(
        epochs=epochs,
        learning_rate=lr,
        batch_size=batch,
        device=device,
        model=model,
        output=None,
    )

    # convert: swap `FakeQuantizedLinear` -> `torch.nn.Linear`, then quantize using `base_config`
    quantize_(model, QATConfig(step=QATStep.CONVERT))

    # save model weights after converting
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    _save_model(model, output, timestamp, 0)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="Quantize",
        description="Static Post Training Quantization with torchao",
    )

    parser.add_argument(
        "-o",
        "--output",
        type=str,
        default="./data/models",
    )

    parser.add_argument("-e", "--epochs", type=int, default=50)
    parser.add_argument("-l", "--lr", type=float, default=0.001)
    parser.add_argument("-b", "--batch", type=int, default=32)
    parser.add_argument(
        "-d",
        "--device",
        type=str,
        choices=list(Device.__args__),
        default="cpu",
        help="Device to use for prediction",
    )

    args = parser.parse_args()

    qat(
        model=MNISTModel(),
        output=args.output,
        epochs=args.epochs,
        lr=args.lr,
        batch=args.batch,
        device=args.device,
    )