add inference time

micromind-toolkit · Nov 27, 2023 · 4e32723 · 4e32723
1 parent d6ae1cd
commit 4e32723
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/recipes/image_classification/inference.py b/recipes/image_classification/inference.py
@@ -14,6 +14,7 @@
 """
 
 import sys
+import time
 import torch
 from train import ImageClassification
 from micromind.utils import parse_configuration
@@ -96,7 +97,9 @@ def acc(pred, batch):
     )
 
     img = preprocess(img.float() / 255)
+    now = time.time()
     logits = mind((img[None],))
+    print("Inference took %.3f ms." % (time.time() - now) * 1e-3)
 
     print(
         "Model prediction: %d with probability: %.2f."