Cleanup mlperf (#797)

* improve factorization

* cleanups

examples/mlperf/model_eval.py

@@ -1,12 +1,9 @@
 import time
 import numpy as np
 from tinygrad.tensor import Tensor
+from tinygrad.helpers import getenv
 
-if __name__ == "__main__":
-  # inference only
-  Tensor.training = False
-  Tensor.no_grad = True
-
+def eval_resnet():
   # Resnet50-v1.5
   from tinygrad.jit import TinyJit
   from models.resnet import ResNet50
@@ -43,6 +40,7 @@ if __name__ == "__main__":
     print(f"****** {n}/{d}  {n*100.0/d:.2f}%")
     st = time.perf_counter()
 
+def eval_rnnt():
   # RNN-T
   from models.rnnt import RNNT
   mdl = RNNT()
@@ -70,3 +68,15 @@ if __name__ == "__main__":
     c += len(tt)
     print(f"WER: {scores/words}, {words} words, raw scores: {scores}, c: {c}")
     st = time.perf_counter()
+
+if __name__ == "__main__":
+  # inference only
+  Tensor.training = False
+  Tensor.no_grad = True
+
+  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
+  for m in models:
+    nm = f"eval_{m}"
+    if nm in globals():
+      print(f"eval {m}")
+      globals()[nm]()

examples/mlperf/model_spec.py

@@ -1,6 +1,6 @@
 # load each model here, quick benchmark
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import GlobalCounters
+from tinygrad.helpers import GlobalCounters, getenv
 
 def test_model(model, *inputs):
   GlobalCounters.reset()
@@ -8,27 +8,25 @@ def test_model(model, *inputs):
   # TODO: return event future to still get the time_sum_s without DEBUG=2
   print(f"{GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.time_sum_s*1000:.2f} ms")
 
-if __name__ == "__main__":
-  # inference only for now
-  Tensor.training = False
-  Tensor.no_grad = True
-
+def spec_resnet():
   # Resnet50-v1.5
   from models.resnet import ResNet50
   mdl = ResNet50()
   img = Tensor.randn(1, 3, 224, 224)
   test_model(mdl, img)
 
-  # Retinanet
+def spec_retinanet():
+  # TODO: Retinanet
+  pass
 
+def spec_unet3d():
   # 3D UNET
   from models.unet3d import UNet3D
   mdl = UNet3D()
-  #mdl.load_from_pretrained()
   img = Tensor.randn(1, 1, 5, 224, 224)
   test_model(mdl, img)
 
-  # RNNT
+def spec_rnnt():
   from models.rnnt import RNNT
   mdl = RNNT()
   mdl.load_from_pretrained()
@@ -36,4 +34,19 @@ if __name__ == "__main__":
   y = Tensor.randn(1, 220)
   test_model(mdl, x, y)
 
-  # BERT-large
+def spec_bert():
+  # TODO: BERT-large
+  pass
+
+if __name__ == "__main__":
+  # inference only for now
+  Tensor.training = False
+  Tensor.no_grad = True
+
+  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(","):
+    nm = f"spec_{m}"
+    if nm in globals():
+      print(f"testing {m}")
+      globals()[nm]()
+
+