from __future__ import annotations import argparse import os import re import sys from pathlib import Path from typing import Any, Dict, Tuple import torch THIS_DIR = Path(__file__).resolve().parent REPO_ROOT = THIS_DIR.parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from quantization.escnet.data.calib import build_calib_loader from quantization.escnet.methods.registry import METHODS from quantization.escnet.quant.calibrate import run_calibration from quantization.escnet.quant.config import QuantConfig from quantization.escnet.quant.prepare import finalize_act_calibration from quantization.escnet.utils.io import save_quant_checkpoint from quantization.escnet.utils.seed import set_seed def _import_load_config(): from model_sources.escnet.ESCNet_model.config import load_config return load_config def _import_escnet(): try: from model_sources.escnet.ESCNet_model.models.ESCNet import ESCNet # type: ignore return ESCNet except Exception as e: raise RuntimeError( "Failed to import ESCNet. Run this from repo root and ensure ./model_sources/escnet/ESCNet_model exists." ) from e def _sanitize_state_dict(sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: # handle common wrappers if "state_dict" in sd and isinstance(sd["state_dict"], dict): sd = sd["state_dict"] if "model" in sd and isinstance(sd["model"], dict): sd = sd["model"] # strip 'module.' from DataParallel keys = list(sd.keys()) if keys and keys[0].startswith("module."): sd = {k[len("module."):]: v for k, v in sd.items()} return sd def _none_if_negative(v: float) -> float | None: return None if float(v) < 0 else float(v) def _apply_recipe(args: argparse.Namespace) -> None: recipe = str(getattr(args, "recipe", "custom")).lower() if recipe == "legacy_our": args.method = "ptq_qdq" args.w_bits = 4 args.a_bits = 4 args.w_mode = "per_channel" args.a_mode = "per_token_group" args.group_size = 32 args.dstg_clip_percentile = -1.0 args.dcrp_tau = 1.0 args.dcrp_zero_ratio = 0.02 args.ep_residual_topk = 0 args.calib_batches = 32 args.calib_batch_size = 4 def _legacy_escnet_qcfg_dict(qcfg: QuantConfig) -> Dict[str, Any]: return { "w_bits": int(qcfg.w_bits), "a_bits": int(qcfg.a_bits), "w_mode": "per_channel", "a_mode": "per_token_group", "group_size": 32, "token_clip_percentile": None, "star_tau": 1.0, "star_zero_ratio": 0.02, "ep_residual_topk": 0, "calib_batches": 32, "calib_batch_size": 4, "exclude_name_regex": None, "verbose": False, } def main() -> None: parser = argparse.ArgumentParser("ESCNet PTQ export (QDQ)") parser.add_argument("--config", type=str, default="configs/escnet/main/config.yaml") parser.add_argument("--ckpt", type=str, required=True, help="Float checkpoint (.pth)") parser.add_argument("--out", type=str, required=True, help="Output quant checkpoint (.pt)") parser.add_argument("--method", type=str, default="ptq_qdq", choices=sorted(METHODS.keys())) parser.add_argument( "--recipe", type=str, default="custom", choices=["custom", "legacy_our"], help="legacy_our uses the released COD-TDQ W4A4 settings.", ) # quant settings parser.add_argument("--w_bits", type=int, default=4, choices=[2, 3, 4, 8]) parser.add_argument("--a_bits", type=int, default=4, choices=[2, 3, 4, 8]) parser.add_argument("--a_mode", type=str, default="per_dstg_group", choices=["per_tensor", "per_channel", "per_dstg_group", "per_token_group"]) parser.add_argument("--w_mode", type=str, default="per_channel", choices=["per_tensor", "per_channel"]) parser.add_argument("--group_size", type=int, default=32) parser.add_argument("--dstg_clip_percentile", "--token_clip_percentile", dest="dstg_clip_percentile", type=float, default=-1.0, help="e.g. 0.99; <0 means None (absmax).") # DCRP parser.add_argument("--dcrp_tau", "--star_tau", dest="dcrp_tau", type=float, default=-1.0, help="DCRP tau; <0 disables") parser.add_argument("--dcrp_zero_ratio", "--star_zero_ratio", dest="dcrp_zero_ratio", type=float, default=-1.0, help="DCRP zero_ratio; <0 disables") # EP residual parser.add_argument("--ep_residual_topk", type=int, default=0) # calibration parser.add_argument("--calib_dir", type=str, default="", help="Dataset dir containing Imgs/GT. Empty->random.") parser.add_argument("--calib_batches", type=int, default=32) parser.add_argument("--calib_batch_size", type=int, default=4) parser.add_argument("--calib_samples", type=int, default=128, help="Used only for random fallback.") parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", type=str, default="cuda:0") parser.add_argument("--exclude_name_regex", type=str, default="", help="Regex to skip wrapping certain modules.") parser.add_argument("--verbose", action="store_true") parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() _apply_recipe(args) load_config = _import_load_config() ESCNet = _import_escnet() cfg = load_config(args.config) set_seed(args.seed if args.seed is not None else getattr(cfg, "rand_seed", 42)) device = torch.device(args.device if torch.cuda.is_available() else "cpu") qcfg = QuantConfig( w_bits=args.w_bits, a_bits=args.a_bits, w_mode=args.w_mode, a_mode=args.a_mode, # type: ignore group_size=args.group_size, dstg_group_size=args.group_size, dstg_clip_percentile=None if args.dstg_clip_percentile < 0 else float(args.dstg_clip_percentile), token_clip_percentile=None if args.dstg_clip_percentile < 0 else float(args.dstg_clip_percentile), dcrp_tau=_none_if_negative(args.dcrp_tau), dcrp_zero_ratio=_none_if_negative(args.dcrp_zero_ratio), star_tau=_none_if_negative(args.dcrp_tau), star_zero_ratio=_none_if_negative(args.dcrp_zero_ratio), ep_residual_topk=int(args.ep_residual_topk), calib_batches=int(args.calib_batches), calib_batch_size=int(args.calib_batch_size), exclude_name_regex=args.exclude_name_regex or None, verbose=bool(args.verbose), ) # build float model (do not autoload backbone weights; we load full ckpt) model = ESCNet(cfg, pretrained=False) model.eval() ckpt = torch.load(args.ckpt, map_location="cpu") ckpt = _sanitize_state_dict(ckpt) missing, unexpected = model.load_state_dict(ckpt, strict=False) if args.verbose: print(f"[load] missing={len(missing)}, unexpected={len(unexpected)}") # wrap preparer = METHODS[args.method] report = preparer(model, qcfg) if args.verbose: print(f"[wrap] {report}") model.to(device) # build loader calib_dir = args.calib_dir if args.calib_dir else None loader = build_calib_loader( config_obj=cfg, calib_dir=calib_dir, img_size=int(getattr(cfg, "img_size", 416)), batch_size=qcfg.calib_batch_size, num_workers=int(args.num_workers), num_samples=int(args.calib_samples), seed=args.seed, ) # calibration + finalize (static modes) run_calibration(model, loader, device=device, num_batches=qcfg.calib_batches) finalize_act_calibration(model, device=device) meta = { "float_ckpt": os.path.abspath(args.ckpt), "config": os.path.abspath(args.config), "method": args.method, "device": str(device), } qcfg_override = _legacy_escnet_qcfg_dict(qcfg) if str(args.recipe).lower() == "legacy_our" else None save_quant_checkpoint(args.out, model.state_dict(), qcfg, meta, qcfg_override=qcfg_override) print(f"[done] saved quant checkpoint -> {args.out}") if __name__ == "__main__": main()