PyTorch QAT INT8 변환 후 정확도 하락 잡기

import torch
import torch.ao.quantization as tq

@torch.no_grad()
def evaluate(model, dataloader, device="cpu"):
    model.eval()
    correct = 0
    total = 0
    for x, y in dataloader:
        x, y = x.to(device), y.to(device)
        out = model(x)
        pred = out.argmax(dim=1)
        correct += (pred == y).sum().item()
        total += y.numel()
    return correct / max(total, 1)

# 1) FP32
acc_fp32 = evaluate(model_fp32, val_loader)

# 2) FakeQuant(QAT 준비 후)
acc_fake = evaluate(model_qat_prepared, val_loader)

# 3) INT8(convert 후)
acc_int8 = evaluate(model_int8, val_loader)

print({"fp32": acc_fp32, "fake": acc_fake, "int8": acc_int8})

import copy
import torch
import torch.ao.quantization as tq

# 백엔드 선택: x86 서버면 보통 fbgemm
torch.backends.quantized.engine = "fbgemm"

model = copy.deepcopy(model_fp32)
model.eval()

# 예: Conv+BN+ReLU fuse (모델 구조에 맞게 수정)
# tq.fuse_modules는 모듈 경로 문자열 리스트를 받음
model = tq.fuse_modules(model, [["conv", "bn", "relu"]], inplace=False)

# QAT 준비
model.train()
model.qconfig = tq.get_default_qat_qconfig("fbgemm")
model_prepared = tq.prepare_qat(model, inplace=False)

# QAT 파인튜닝(짧게라도 필수)
optimizer = torch.optim.SGD(model_prepared.parameters(), lr=1e-4, momentum=0.9)
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(3):
    model_prepared.train()
    for x, y in train_loader:
        optimizer.zero_grad()
        out = model_prepared(x)
        loss = criterion(out, y)
        loss.backward()
        optimizer.step()

# 변환
model_prepared.eval()
model_int8 = tq.convert(model_prepared, inplace=False)

import torch.ao.quantization as tq

activation_observer = tq.HistogramObserver.with_args(
    reduce_range=False
)

weight_observer = tq.default_per_channel_weight_observer

qconfig = tq.QConfig(
    activation=activation_observer,
    weight=weight_observer
)

model.qconfig = qconfig
model_prepared = tq.prepare_qat(model, inplace=False)

import torch.ao.quantization as tq

def set_qat_freeze(model):
    # 옵저버 비활성화
    model.apply(tq.disable_observer)
    # FakeQuant 고정(스케일/제로포인트 업데이트 중지)
    model.apply(tq.disable_fake_quant)

# 예: 전체 epoch 중 후반부에만 freeze
for epoch in range(10):
    train_one_epoch(model_prepared)
    if epoch == 7:
        model_prepared.apply(tq.disable_observer)
    if epoch == 8:
        model_prepared.apply(tq.disable_fake_quant)

import torch.ao.quantization as tq

# 예: 마지막 분류기 레이어는 FP32 유지
model.classifier.qconfig = None

# 또는 특정 서브모듈 전체를 제외
model.head.qconfig = None

model_prepared = tq.prepare_qat(model, inplace=False)

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

class QuantWrapper(nn.Module):
    def __init__(self, core):
        super().__init__()
        self.quant = tq.QuantStub()
        self.core = core
        self.dequant = tq.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.core(x)
        x = self.dequant(x)
        return x

model = QuantWrapper(model_fp32)

import torch

torch.backends.quantized.engine = "fbgemm"  # 또는 "qnnpack"
print(torch.backends.quantized.engine)

def assert_input_ok(x):
    assert x.dtype == torch.float32
    assert x.ndim == 4
    # 값 범위 점검(예: 0-1 가정)
    assert x.min().item() >= -0.1 and x.max().item() <= 1.1

import torch

def collect_activations(model, layer_names):
    acts = {}
    hooks = []

    name_to_module = dict(model.named_modules())

    def make_hook(name):
        def hook(m, inp, out):
            # out이 tuple일 수 있어 방어
            t = out[0] if isinstance(out, (tuple, list)) else out
            acts[name] = t.detach().float().cpu()
        return hook

    for n in layer_names:
        hooks.append(name_to_module[n].register_forward_hook(make_hook(n)))

    return acts, hooks

@torch.no_grad()
def compare_one_batch(model_a, model_b, x, layer_names):
    acts_a, hooks_a = collect_activations(model_a, layer_names)
    acts_b, hooks_b = collect_activations(model_b, layer_names)

    _ = model_a(x)
    _ = model_b(x)

    for h in hooks_a + hooks_b:
        h.remove()

    diffs = {}
    for n in layer_names:
        a = acts_a[n]
        b = acts_b[n]
        diffs[n] = (a - b).abs().mean().item()
    return diffs

# 사용 예: 중요한 블록 몇 개만 찍어도 방향이 보임
layer_names = [
    "core.layer1.0.conv1",
    "core.layer2.0.conv1",
    "core.layer4.1.conv2",
]

diffs = compare_one_batch(model_fp32.eval(), model_prepared.eval(), x_batch, layer_names)
print(diffs)