Spaces:
Running on Zero
Running on Zero
| 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() | |