import torch, torch.nn as nn
import torch.ao.quantization as tq

class QNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = tq.QuantStub()                        # FP32 -> INT8 ở đầu vào
        self.conv1, self.bn1, self.relu1 = nn.Conv2d(1, 16, 3, padding=1, bias=False), nn.BatchNorm2d(16), nn.ReLU()
        self.dw, self.bn2, self.relu2 = nn.Conv2d(16, 16, 3, padding=1, groups=16, bias=False), nn.BatchNorm2d(16), nn.ReLU()
        self.pw, self.bn3, self.relu3 = nn.Conv2d(16, 32, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU()
        self.pool = nn.MaxPool2d(2)
        self.conv4, self.bn4, self.relu4 = nn.Conv2d(32, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.ReLU()
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(32, 10)
        self.dequant = tq.DeQuantStub()                    # INT8 -> FP32 ở đầu ra

    def forward(self, x):
        x = self.quant(x)
        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.dw(x)))
        x = self.pool(self.relu3(self.bn3(self.pw(x))))
        x = self.relu4(self.bn4(self.conv4(x)))
        x = self.fc(torch.flatten(self.gap(x), 1))
        return self.dequant(x)

model = QNet(); model.load_state_dict(fp32_state_dict, strict=False); model.eval()

torch.backends.quantized.engine = "x86"                  # ARM / điện thoại: "qnnpack"
model.qconfig = tq.get_default_qconfig("x86")            # trọng số per-channel đối xứng, activation per-tensor
tq.fuse_modules(model, [["conv1", "bn1", "relu1"], ["dw", "bn2", "relu2"],
                        ["pw", "bn3", "relu3"], ["conv4", "bn4", "relu4"]], inplace=True)   # Mục 7.3–7.4
tq.prepare(model, inplace=True)                          # chèn observer
with torch.no_grad():
    model(X_calib)                                       # calibration (Chương 8)
tq.convert(model, inplace=True)                          # thay bằng kernel INT8 thật
print(model.conv1)

import torch.ao.nn.intrinsic.qat as nniqat

model = QNet(); model.load_state_dict(fp32_state_dict, strict=False)   # BẮT ĐẦU TỪ MÔ HÌNH ĐÃ HUẤN LUYỆN
model.train()
model.qconfig = tq.get_default_qat_qconfig("x86")
tq.fuse_modules_qat(model, [["conv1", "bn1", "relu1"], ["dw", "bn2", "relu2"],
                            ["pw", "bn3", "relu3"], ["conv4", "bn4", "relu4"]], inplace=True)
tq.prepare_qat(model, inplace=True)                      # chèn FakeQuantize, mô phỏng gộp BN

opt = torch.optim.Adam(model.parameters(), lr=5e-4)     # learning rate nhỏ: đây là fine-tune
for epoch in range(10):
    if epoch == 6: model.apply(tq.disable_observer)      # cố định dải lượng tử
    if epoch == 7: model.apply(nniqat.freeze_bn_stats)   # cố định thống kê BatchNorm
    train_one_epoch(model, opt)

model.eval()
qmodel = tq.convert(model)

from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e
import torchao.quantization.pt2e.quantizer.x86_inductor_quantizer as xiq

m = fp32_model.eval()                                        # mô hình gốc, viết bằng F.relu cũng được
exported = torch.export.export(
    m, (X_calib[:4],),
    dynamic_shapes={"x": {0: torch.export.Dim("batch")}}     # cho phép batch size thay đổi
).module()

quantizer = xiq.X86InductorQuantizer().set_global(xiq.get_default_x86_inductor_quantization_config())
prepared = prepare_pt2e(exported, quantizer)                 # chèn observer
with torch.no_grad():
    prepared(X_calib)                                        # calibration
quantized = convert_pt2e(prepared)                           # đồ thị có các cặp quantize/dequantize
# Để có kernel INT8 nhanh: torch.compile(quantized)

import torch.ao.quantization as tq
from torchao.quantization import quantize_, Int8WeightOnlyConfig, Int8DynamicActivationInt8WeightConfig

m1 = tq.quantize_dynamic(copy.deepcopy(base), {nn.Linear}, dtype=torch.qint8)     # API cũ
m2 = copy.deepcopy(base); quantize_(m2, Int8WeightOnlyConfig())                   # torchao
m3 = copy.deepcopy(base); quantize_(m3, Int8DynamicActivationInt8WeightConfig())  # torchao

import tensorflow as tf

def make_converter():                       # tạo converter MỚI cho mỗi cấu hình để tránh cấu hình cũ dính lại
    return tf.lite.TFLiteConverter.from_keras_model(model)

def representative_data_gen():              # dùng cho calibration (Chương 8)
    for x in tf.data.Dataset.from_tensor_slices(calib_images).batch(1).take(300):
        yield [x]

# 1) Dynamic range: trọng số INT8 lượng tử sẵn; activation lượng tử lúc chạy ở các op hỗ trợ
c = make_converter()
c.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_dynamic = c.convert()

# 2) Float16: trọng số lưu FP16
c = make_converter()
c.optimizations = [tf.lite.Optimize.DEFAULT]
c.target_spec.supported_types = [tf.float16]
tflite_fp16 = c.convert()

# 3) Full integer INT8: bắt buộc có representative dataset
c = make_converter()
c.optimizations = [tf.lite.Optimize.DEFAULT]
c.representative_dataset = representative_data_gen
c.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]   # op không có bản INT8 -> báo lỗi, không lặng lẽ chạy FP32
c.inference_input_type = tf.int8
c.inference_output_type = tf.int8
tflite_int8 = c.convert()

# 4) 16x8: activation INT16, trọng số INT8; cho mô hình nhạy với sai số activation
c = make_converter()
c.optimizations = [tf.lite.Optimize.DEFAULT]
c.representative_dataset = representative_data_gen
c.target_spec.supported_ops = [tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8]
tflite_16x8 = c.convert()

interpreter = tf.lite.Interpreter(model_content=tflite_int8)
interpreter.allocate_tensors()
inp, out = interpreter.get_input_details()[0], interpreter.get_output_details()[0]
S, Z = inp["quantization"]
x_q = np.clip(np.round(x / S + Z), -128, 127).astype(np.int8)            # công thức Mục 3.2
interpreter.set_tensor(inp["index"], x_q[None])
interpreter.invoke()
y_q = interpreter.get_tensor(out["index"])
S_o, Z_o = out["quantization"]
y = S_o * (y_q.astype(np.float32) - Z_o)                                  # dequantize

import tensorflow_model_optimization as tfmot

q_aware = tfmot.quantization.keras.quantize_model(model)   # chèn fake-quant (Hình 10)
q_aware.compile(optimizer=tf.keras.optimizers.Adam(1e-4),
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
                metrics=["accuracy"])
q_aware.fit(train_images, train_labels, epochs=2, validation_split=0.1)

c = tf.lite.TFLiteConverter.from_keras_model(q_aware)
c.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_qat = c.convert()

import os, time, torch

def size_mb(model):
    torch.save(model.state_dict(), "/tmp/_m.pt")
    s = os.path.getsize("/tmp/_m.pt") / 2**20
    os.remove("/tmp/_m.pt")
    return s

@torch.no_grad()
def latency_ms(model, x, n=100, warmup=10):
    for _ in range(warmup):
        model(x)
    t = time.perf_counter()
    for _ in range(n):
        model(x)
    return (time.perf_counter() - t) / n * 1000