src/quant/quantize_yolo.py

"""
YOLOv8s 검출기 양자화 스크립트.

Phase 1의 base_W8A8.py / multimodal_w8a8_smoothquant.py 패턴을 YOLOv8에 포팅.
PyTorch 레벨 fake-quantization 후 ONNX 내보내기.

지원 모드:
  w8a8     — W8A8 PTQ (per-channel MinMax, Phase 1 base_W8A8 동일 방식)
  w4a16    — W4A16 PTQ (4-bit 가중치 / 16-bit 활성화)
  smoothquant — SmoothQuant + W8A8 (Phase 1 multimodal_smoothquant 동일 방식)

사용법:
  python src/quant/quantize_yolo.py --mode w8a8
  python src/quant/quantize_yolo.py --mode w4a16
  python src/quant/quantize_yolo.py --mode smoothquant --calib_batches 10
"""

import argparse
import shutil
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

ROOT = Path(__file__).parent.parent.parent
MODEL_SPACE = ROOT / "model_space"
WEIGHTS = ROOT / "runs" / "detect" / "edge_sign_v2_e0_full3" / "weights" / "best.pt"
DATA_DIR = ROOT / "data" / "yolo_signs"
YOLO_DATASET = DATA_DIR / "dataset.yaml"

MODEL_SPACE.mkdir(parents=True, exist_ok=True)

# ────────────────────────────────────────────
# 공통 유틸
# ────────────────────────────────────────────


def load_yolo_model(weights=WEIGHTS):
    from ultralytics import YOLO

    return YOLO(str(weights))


def _is_quantizable(name: str, module: nn.Module) -> bool:
    """Conv2d / Linear 중 Detection Head 제외."""
    if not isinstance(module, (nn.Conv2d, nn.Linear)):
        return False
    # YOLO detection head (model.22.*) — 마지막 출력 보호
    skip_keywords = ["dfl", "detect"]
    return not any(k in name.lower() for k in skip_keywords)


def export_to_onnx(yolo_model, out_name: str, opset: int = 14) -> Path:
    """수정된 PyTorch 모델을 ONNX로 내보내기 (ultralytics .export() 사용)."""
    result = yolo_model.export(
        format="onnx",
        imgsz=640,
        half=False,
        simplify=True,
        opset=opset,
        dynamic=False,
    )
    src = Path(result)
    dst = MODEL_SPACE / out_name
    shutil.copy2(src, dst)
    size_mb = dst.stat().st_size / 1024 / 1024
    print(f"  → 저장: {dst} ({size_mb:.2f} MB)")
    return dst


def export_nn_to_onnx(nn_model: nn.Module, out_name: str, opset: int = 14) -> Path:
    """
    이미 수정된 nn.Module을 torch.onnx.export로 직접 내보내기.
    SmoothQuant처럼 wrapper가 포함된 경우 사용 (ultralytics .export()의 fuse() 충돌 회피).
    """
    import onnx
    import onnxslim

    MODEL_SPACE.mkdir(parents=True, exist_ok=True)
    dst = MODEL_SPACE / out_name
    tmp = MODEL_SPACE / ("_tmp_" + out_name)

    nn_model.eval()
    dummy = torch.randn(1, 3, 640, 640)
    with torch.no_grad():
        torch.onnx.export(
            nn_model,
            dummy,
            str(tmp),
            opset_version=opset,
            input_names=["images"],
            output_names=["output0"],
            do_constant_folding=True,
            dynamo=False,  # TorchScript 기반 exporter 사용 (PyTorch 2.x 호환)
        )

    # onnxslim으로 최적화
    try:
        slimmed = onnxslim.slim(str(tmp))
        onnx.save(slimmed, str(dst))
        tmp.unlink(missing_ok=True)
    except Exception:
        tmp.rename(dst)

    size_mb = dst.stat().st_size / 1024 / 1024
    print(f"  → 저장: {dst} ({size_mb:.2f} MB)")
    return dst


def verify_onnx(path: Path):
    import onnxruntime as ort

    sess = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    dummy = np.random.randn(1, 3, 640, 640).astype(np.float32)
    out = sess.run(None, {sess.get_inputs()[0].name: dummy})
    print(f"  검증 OK: output shape = {out[0].shape}")


# ────────────────────────────────────────────
# W8A8 PTQ (Phase 1 base_W8A8.py 동일 방식)
# ────────────────────────────────────────────


def apply_w8a8_ptq(model_nn: nn.Module) -> int:
    """
    Conv2d / Linear 레이어에 per-channel MinMax W8A8 fake-quantization 적용.
    Phase 1의 apply_w8a8_ptq() 와 동일한 로직.
    """
    quantized = 0
    for name, module in model_nn.named_modules():
        if not _is_quantizable(name, module):
            continue
        with torch.no_grad():
            w = module.weight.data
            # Per-output-channel MinMax scale
            if w.dim() == 4:  # Conv2d: [out, in, kH, kW]
                max_val = w.view(w.size(0), -1).abs().max(dim=1)[0].view(-1, 1, 1, 1)
            else:  # Linear: [out, in]
                max_val = w.abs().max(dim=1)[0].view(-1, 1)
            scale = (max_val / 127.0).clamp(min=1e-8)
            q_w = torch.round(w / scale).clamp(-128, 127)
            module.weight.data = q_w * scale  # fake-dequant
        quantized += 1
    return quantized


def run_w8a8(weights=WEIGHTS):
    print("\n[W8A8 PTQ] 시작")
    yolo = load_yolo_model(weights)
    nn_model = yolo.model

    n = apply_w8a8_ptq(nn_model)
    print(f"  양자화 레이어: {n}개 (Detection Head 제외)")

    out = export_to_onnx(yolo, "yolov8s_signs_w8a8.onnx")
    verify_onnx(out)
    return out


# ────────────────────────────────────────────
# W4A16 PTQ (4-bit 가중치 / FP16 활성화)
# ────────────────────────────────────────────


def apply_w4a16_ptq(model_nn: nn.Module) -> int:
    """
    4-bit 가중치 양자화 시뮬레이션 (활성화는 FP32 유지).
    Phase 1의 W4A16 QAT 가중치 표현과 동일한 INT4 범위(-8 ~ 7).
    """
    quantized = 0
    for name, module in model_nn.named_modules():
        if not _is_quantizable(name, module):
            continue
        with torch.no_grad():
            w = module.weight.data
            if w.dim() == 4:
                max_val = w.view(w.size(0), -1).abs().max(dim=1)[0].view(-1, 1, 1, 1)
            else:
                max_val = w.abs().max(dim=1)[0].view(-1, 1)
            scale = (max_val / 7.0).clamp(min=1e-8)  # INT4: [-8, 7]
            q_w = torch.round(w / scale).clamp(-8, 7)
            module.weight.data = q_w * scale
        quantized += 1
    return quantized


def run_w4a16(weights=WEIGHTS):
    print("\n[W4A16 PTQ] 시작")
    yolo = load_yolo_model(weights)
    nn_model = yolo.model

    n = apply_w4a16_ptq(nn_model)
    print(f"  양자화 레이어: {n}개 (4-bit 가중치)")

    out = export_to_onnx(yolo, "yolov8s_signs_w4a16.onnx")
    verify_onnx(out)
    return out


# ────────────────────────────────────────────
# SmoothQuant + W8A8 (Phase 1 동일 방식)
# ────────────────────────────────────────────


def _build_calib_loader(num_batches=10, batch_size=4):
    """val 이미지를 캘리브레이션 데이터로 사용."""
    import cv2
    from torch.utils.data import DataLoader, Dataset

    img_dir = DATA_DIR / "images" / "val"
    img_paths = sorted(img_dir.rglob("*.jpg"))[: num_batches * batch_size]

    class YOLOImageDataset(Dataset):
        def __init__(self, paths, imgsz=640):
            self.paths = paths
            self.imgsz = imgsz

        def __len__(self):
            return len(self.paths)

        def __getitem__(self, idx):
            img = cv2.imread(str(self.paths[idx]))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (self.imgsz, self.imgsz))
            tensor = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
            return tensor

    ds = YOLOImageDataset(img_paths)
    return DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=0)


class _SmoothWrapper(nn.Module):
    """
    SmoothQuant Wrapper: forward에서 입력을 1/s로 스케일링 후 가중치(s 흡수+W8) 레이어 실행.
    ONNX export 시 스케일 나눗셈이 그래프에 포함됨 (Phase 1 SmoothQuantWrapper 동일 방식).
    """

    def __init__(self, module: nn.Module, smooth_scale: torch.Tensor):
        super().__init__()
        self.module = module
        self.register_buffer("smooth_scale", smooth_scale)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.dim() == 4:
            x = x / self.smooth_scale.view(1, -1, 1, 1)
        else:
            x = x / self.smooth_scale
        return self.module(x)


def apply_smoothquant(model_nn: nn.Module, calib_loader, alpha: float = 0.5) -> int:
    """
    SmoothQuant: 활성화 캘리브레이션 → wrapper로 입력 스케일링 + 가중치 흡수 + W8A8.
    wrapper.forward()에 x/s 연산이 포함돼 ONNX export 시 그래프에 반영됨.
    """
    device = next(model_nn.parameters()).device
    model_nn.eval()

    # 1. 활성화 최대값 수집 (per-input-channel)
    act_max: dict = {}
    hooks = []

    def make_hook(name):
        def hook(module, inp, out):
            x = inp[0].detach().abs()
            if x.dim() == 4:
                ch_max = x.amax(dim=(0, 2, 3))
            else:
                ch_max = x.amax(dim=0) if x.dim() >= 2 else x
            act_max[name] = torch.max(act_max[name], ch_max) if name in act_max else ch_max

        return hook

    target_names = [name for name, m in model_nn.named_modules() if _is_quantizable(name, m)]
    target_mods = dict(model_nn.named_modules())

    for name in target_names:
        hooks.append(target_mods[name].register_forward_hook(make_hook(name)))

    print(f"  캘리브레이션 중 ({len(calib_loader)} 배치)...")
    with torch.no_grad():
        for batch in calib_loader:
            model_nn(batch.to(device))
    for h in hooks:
        h.remove()

    # 2. Wrapper 교체 + 가중치 W8A8 적용
    def _set_module(root, dotted_name, new_module):
        parts = dotted_name.split(".")
        parent = root
        for p in parts[:-1]:
            parent = getattr(parent, p)
        setattr(parent, parts[-1], new_module)

    quantized = 0
    for name in target_names:
        if name not in act_max:
            continue
        module = target_mods[name]
        w = module.weight.data
        a_max = act_max[name].to(device).clamp(min=1e-8)

        if w.dim() == 4:
            in_ch = w.size(1)
        else:
            in_ch = w.size(1)

        # a_max 채널 수 맞추기
        if a_max.shape[0] != in_ch:
            if a_max.shape[0] > in_ch:
                a_max = a_max[:in_ch]
            else:
                pad = a_max.mean().expand(in_ch - a_max.shape[0])
                a_max = torch.cat([a_max, pad])

        # per-input-channel weight max
        if w.dim() == 4:
            w_max = w.abs().amax(dim=(0, 2, 3)).clamp(min=1e-8)
        else:
            w_max = w.abs().amax(dim=0).clamp(min=1e-8)

        smooth_s = (a_max**alpha) / (w_max ** (1 - alpha) + 1e-8)
        smooth_s = smooth_s.clamp(1e-3, 1e3)

        # 가중치에 smooth_s 흡수 + W8 fake-quant
        with torch.no_grad():
            if w.dim() == 4:
                w_scaled = w * smooth_s.view(1, in_ch, 1, 1)
                out_max = w_scaled.view(w.size(0), -1).abs().max(dim=1)[0].view(-1, 1, 1, 1)
            else:
                w_scaled = w * smooth_s.view(1, in_ch)
                out_max = w_scaled.abs().max(dim=1)[0].view(-1, 1)
            q_scale = (out_max / 127.0).clamp(min=1e-8)
            q_w = torch.round(w_scaled / q_scale).clamp(-128, 127)
            module.weight.data = q_w * q_scale

        # Wrapper 교체 (ONNX export 시 x/smooth_s 연산 포함)
        wrapper = _SmoothWrapper(module, smooth_s)
        _set_module(model_nn, name, wrapper)
        quantized += 1

    return quantized


def run_smoothquant(weights=WEIGHTS, calib_batches=10, alpha=0.5):
    print("\n[SmoothQuant + W8A8] 시작")
    yolo = load_yolo_model(weights)
    nn_model = yolo.model

    # ★ fuse() 먼저: Conv+BN 융합 후 SmoothWrapper 교체
    #   그래야 ultralytics fuse() 재호출 없이 torch.onnx.export 가능
    nn_model = nn_model.fuse()
    nn_model.eval()

    calib_loader = _build_calib_loader(num_batches=calib_batches)
    n = apply_smoothquant(nn_model, calib_loader, alpha=alpha)
    print(f"  SmoothQuant 적용 레이어: {n}개 (alpha={alpha})")

    # torch.onnx.export 직접 사용 (ultralytics .export()의 fuse() 재호출 회피)
    out = export_nn_to_onnx(nn_model, "yolov8s_signs_smoothquant.onnx")
    verify_onnx(out)
    return out


# ────────────────────────────────────────────
# CLI
# ────────────────────────────────────────────


def main():
    parser = argparse.ArgumentParser(description="YOLOv8s 검출기 양자화")
    parser.add_argument(
        "--mode", choices=["w8a8", "w4a16", "smoothquant", "all"], default="all", help="양자화 모드"
    )
    parser.add_argument("--weights", type=str, default=str(WEIGHTS), help="학습된 best.pt 경로")
    parser.add_argument(
        "--calib_batches", type=int, default=10, help="SmoothQuant 캘리브레이션 배치 수"
    )
    parser.add_argument(
        "--alpha", type=float, default=0.5, help="SmoothQuant alpha (0=weight만, 1=activation만)"
    )
    args = parser.parse_args()

    modes = ["w8a8", "w4a16", "smoothquant"] if args.mode == "all" else [args.mode]

    for mode in modes:
        if mode == "w8a8":
            run_w8a8(args.weights)
        elif mode == "w4a16":
            run_w4a16(args.weights)
        elif mode == "smoothquant":
            run_smoothquant(args.weights, args.calib_batches, args.alpha)

    print("\n모든 양자화 완료. model_space/ 디렉토리 확인:")
    for f in sorted(MODEL_SPACE.glob("yolov8s_signs_*.onnx")):
        print(f"  {f.name}: {f.stat().st_size / 1024 / 1024:.2f} MB")


if __name__ == "__main__":
    main()