| import argparse | |
| import onnx | |
| import torch | |
| from fvcore.nn import FlopCountAnalysis | |
| from onnx2torch import convert | |
| def get_args(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| "--model-ckpt", | |
| type=str, | |
| help="Model checkpoint. Both PyTorch and ONNX models can be used.", | |
| ) | |
| return parser.parse_args() | |
| def main(): | |
| args = get_args() | |
| if args.model_ckpt.endswith(".onnx"): | |
| onnx_model = onnx.load(args.model_ckpt) | |
| model = convert(onnx_model) | |
| elif args.model_ckpt.endswith((".pth", ".pt")): | |
| checkpoint = torch.load(args.model_ckpt, map_location="cpu") | |
| model = checkpoint["model_ckpt"] | |
| if "model_ema" in checkpoint: | |
| state_dict = {} | |
| for key, value in checkpoint["model_ema"].items(): | |
| if not "module." in key: | |
| continue | |
| state_dict[key.replace("module.", "")] = value | |
| model.load_state_dict(state_dict) | |
| else: | |
| raise RuntimeError( | |
| f"Cannot process file {args.model_ckpt} with extension {args.model_ckpt.split('.')[-1]}" | |
| ) | |
| model.eval() | |
| flops = FlopCountAnalysis(model.cpu(), torch.ones((1, 3, 224, 224))) | |
| flops = flops.total() | |
| print(f"MMACs = {flops/1e6}") | |
| if __name__ == "__main__": | |
| main() | |