Quant/main.py at main · ilfpns/Quant · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
from __future__ import annotations

import torch

from quanti.export import default_onnx_filename, export_to_onnx, onnx_file_size_kb
from quanti.model import SimpleMLP
from quanti.ptq import dynamic_quantize, static_quantize
from quanti.qat import build_self_distillation_data, qat_quantize
from quanti.report import (
    SpeedResult,
    max_output_diff,
    measure,
    measure_latency_ms,
    print_accuracy_report,
    print_size_report,
    print_speed_table,
)

INPUT_SIZE = 128
CALIBRATION_BATCHES = 32
QAT_TRAIN_BATCHES = 200
QAT_EPOCHS = 3


def main() -> None:
    torch.manual_seed(0)
    fp32_model = SimpleMLP(input_size=INPUT_SIZE)
    fp32_model.eval()

    sample_input = torch.randn(1, INPUT_SIZE)
    calibration_data = [torch.randn(1, INPUT_SIZE) for _ in range(CALIBRATION_BATCHES)]
    qat_train_data = build_self_distillation_data(fp32_model, INPUT_SIZE, QAT_TRAIN_BATCHES)

    dynamic_model = dynamic_quantize(fp32_model)
    static_model = static_quantize(fp32_model, calibration_data)
    qat_model = qat_quantize(fp32_model, qat_train_data, epochs=QAT_EPOCHS)

    with torch.no_grad():
        fp32_output = fp32_model(sample_input)
        dynamic_output = dynamic_model(sample_input)
        static_output = static_model(sample_input)
        qat_output = qat_model(sample_input)

    print_size_report(
        [
            measure("FP32", fp32_model),
            measure("Dynamic PTQ (INT8)", dynamic_model),
            measure("Static PTQ (INT8)", static_model),
            measure("QAT (INT8)", qat_model),
        ]
    )

    speed_results = [
        SpeedResult("FP32", measure_latency_ms(fp32_model, sample_input)),
        SpeedResult(
            "Dynamic PTQ",
            measure_latency_ms(dynamic_model, sample_input),
            max_output_diff(fp32_output, dynamic_output),
        ),
        SpeedResult(
            "Static PTQ",
            measure_latency_ms(static_model, sample_input),
            max_output_diff(fp32_output, static_output),
        ),
        SpeedResult(
            "QAT",
            measure_latency_ms(qat_model, sample_input),
            max_output_diff(fp32_output, qat_output),
        ),
    ]

    print_speed_table(speed_results)
    print_accuracy_report(speed_results)

    onnx_path = default_onnx_filename(fp32_model)
    export_to_onnx(fp32_model, sample_input, onnx_path)
    print()
    print("===== ONNX Export =====")
    print(f"Path: {onnx_path}")
    print(f"Size: {onnx_file_size_kb(onnx_path):.2f} KB")
    print("========================")


if __name__ == "__main__":
    main()