IQ.Pilot Release Commit @ b6534c0

artifacts/package_sources/tinygrad/examples/test_onnx_imagenet.py

@@ -0,0 +1,82 @@
+import random, sys
+import numpy as np
+from extra.datasets.imagenet import get_imagenet_categories, get_val_files, center_crop
+from examples.benchmark_onnx import load_onnx_model
+from PIL import Image
+from tinygrad import Tensor, dtypes, GlobalCounters
+from tinygrad.helpers import fetch, getenv
+
+# works:
+#  ~70% - https://github.com/onnx/models/raw/refs/heads/main/validated/vision/classification/resnet/model/resnet50-v2-7.onnx
+#  ~43% - https://github.com/onnx/models/raw/refs/heads/main/Computer_Vision/alexnet_Opset16_torch_hub/alexnet_Opset16.onnx
+#  ~72% - https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
+#  ~71% - https://github.com/axinc-ai/onnx-quantization/raw/refs/heads/main/models/mobilenetv2_1.0.opt.onnx
+#  ~67% - https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7-quantized.onnx
+# broken:
+#  https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
+#  https://huggingface.co/qualcomm/MobileNet-v2-Quantized/resolve/main/MobileNet-v2-Quantized.onnx
+#  ~35% - https://github.com/axinc-ai/onnx-quantization/raw/refs/heads/main/models/mobilenev2_quantized.onnx
+
+# QUANT=1 python3 examples/test_onnx_imagenet.py
+#   https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
+# python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+# VIZ=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
+
+def imagenet_dataloader(cnt=0):
+  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
+  input_std = Tensor([0.229, 0.224, 0.225]).reshape(1, -1, 1, 1)
+  files = get_val_files()
+  random.shuffle(files)
+  files = files[:cnt]
+  cir = get_imagenet_categories()
+  for fn in files:
+    img = Image.open(fn)
+    img = img.convert('RGB') if img.mode != "RGB" else img
+    img = center_crop(img)
+    img = np.array(img)
+    img = Tensor(img).permute(2,0,1).reshape(1,3,224,224)
+    img = ((img.cast(dtypes.float32)/255.0) - input_mean) / input_std
+    y = cir[fn.split("/")[-2]]
+    yield img,y
+
+if __name__ == "__main__":
+  fn = sys.argv[1]
+  if getenv("QUANT"):
+    from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantFormat, QuantType, CalibrationDataReader
+    model_fp32 = fetch(fn)
+    fn = '/tmp/model.quant.onnx'
+    if getenv("DYNAMIC"):
+      quantize_dynamic(model_fp32, fn)
+    else:
+      class ImagenetReader(CalibrationDataReader):
+        def __init__(self):
+          self.iter = imagenet_dataloader(cnt=1000)
+        def get_next(self) -> dict:
+          try:
+            img,y = next(self.iter)
+          except StopIteration:
+            return None
+          return {"input": img.numpy()}
+      quantize_static(model_fp32, fn, ImagenetReader(), quant_format=QuantFormat.QDQ, per_channel=False,
+                      activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8,
+                      extra_options={"ActivationSymmetric": False})
+
+  run_onnx_jit, input_specs = load_onnx_model(fetch(fn))
+  t_name, t_spec = list(input_specs.items())[0]
+  assert t_spec.shape[1:] == (3,224,224), f"shape is {t_spec.shape}"
+
+  hit = 0
+  for i,(img,y) in enumerate(imagenet_dataloader(cnt:=getenv("CNT", 100))):
+    GlobalCounters.reset()
+    p = run_onnx_jit(**{t_name:img})
+    assert p.shape == (1,1000)
+    t = p.to('cpu').argmax().item()
+    hit += y==t
+    print(f"target: {y:3d}  pred: {t:3d}  acc: {hit/(i+1)*100:.2f}%")
+
+  MS_TARGET = 13.4
+  print(f"need {GlobalCounters.global_ops/1e9*(1000/MS_TARGET):.2f} GFLOPS for {MS_TARGET:.2f} ms")
+
+  if cnt >= 2:
+    import pickle
+    with open("/tmp/im.pkl", "wb") as f: pickle.dump(run_onnx_jit, f)