| |
| import argparse, os |
| import torch |
| from pointcept.models.builder import build_model |
|
|
| |
| from pointcept.utils.config import get_cfg_from_file |
|
|
| from pointcept.models.utils.quant_0920 import ( |
| apply_quantization_0920, print_avg_bits_0920 |
| ) |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--config-file", required=True) |
| ap.add_argument("--options", nargs=argparse.REMAINDER, default=None) |
| ap.add_argument("--ckpt", default="", help="optional: load checkpoint then wrap quant") |
| ap.add_argument("--wbits", type=int, default=2) |
| ap.add_argument("--abits", type=int, default=8) |
| args = ap.parse_args() |
|
|
| cfg = get_cfg_from_file(args.config_file) |
| if args.options: |
| cfg.merge_from_list(args.options) |
|
|
| model = build_model(cfg.model) |
| if args.ckpt and os.path.isfile(args.ckpt): |
| sd = torch.load(args.ckpt, map_location="cpu") |
| if "state_dict" in sd: |
| sd = sd["state_dict"] |
| model.load_state_dict(sd, strict=False) |
|
|
| model = apply_quantization_0920(model, enable=True, w_bits=args.wbits, a_bits=args.abits, |
| quantize_first_last=False, exclude_name_hints=["cls_head", "embedding.stem"]) |
| print_avg_bits_0920(model) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
