add ViT x4.8 ONNX, small fixes in test.py, and allow measuring macs on ONNX

README.md

@@ -28,7 +28,8 @@ Evaluation code is also based on Torchvision references.
 | Model                    | Latency (MMACs) | Accuracy (%)  |
 |--------------------------|:---------------:|:-------------:|
 | **ViT-B/32 Torchvision** |     4413.99     |     75.91     |
-| **ViT-B/32 ENOT**        | 492.23 (x8.97)  | 73.72 (-2.19) |
+| **ViT-B/32 ENOT (x4.8)** | 911.80 (x4.84)  | 75.68 (-0.23) |
+| **ViT-B/32 ENOT (x9)**   | 490.78 (x8.99)  | 73.72 (-2.19) |
 
 ## MobileNetV2
 

ViT-B-32/{ViT-B-32-ENOT.pth → ViT-B-32-ENOT-x4_8.onnx}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:92a81cef913af4012215215400049317168939624f09d76e9043aee2342af356
-size 157444613
+oid sha256:3597a973923ab2be41e046ef08cbccadd67279853e78185194f906086063e626
+size 72211694

measure_mac.py

@@ -1,12 +1,18 @@
 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)
+    parser.add_argument(
+        "--model-ckpt",
+        type=str,
+        help="Model checkpoint. Both PyTorch and ONNX models can be used.",
+    )
 
     return parser.parse_args()
 
@@ -14,8 +20,24 @@ def get_args():
 def main():
     args = get_args()
 
-    checkpoint = torch.load(args.model_ckpt, map_location="cpu")
-    model = checkpoint["model_ckpt"]
+    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)))

requirements.txt

@@ -3,3 +3,4 @@ torchvision==0.14.1
 fvcore==0.1.5.post20221221
 onnxruntime-gpu==1.15.1
 onnx==1.13.1
+onnx2torch==1.5.6

test.py

@@ -96,6 +96,9 @@ def load_data(valdir):
 def main(args):
     print(args)
 
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
     if torch.cuda.is_available():
         device = torch.device("cuda")
     else:
@@ -128,6 +131,7 @@ def main(args):
                 state_dict[key.replace("module.", "")] = value
             model.load_state_dict(state_dict)
         model = model.to(device)
+        model.eval()
 
     accuracy = evaluate(
         model=model,