modify benchmark to make early stop work

Author: Zhang Yan

dl_bench/mlp.py

@@ -87,7 +87,10 @@ def __init__(self, params) -> None:
 
         name = params.get("name", "size5")
         net = get_mlp(n_chans_in=IN_FEAT, n_chans_out=N_CLASSES, name=name)
+        min_batches = int(params.get("min_batches", 10))
+        min_seconds = int(params.get("min_seconds", 10))
 
         super().__init__(
-            net=net, in_shape=in_shape, dataset=dataset, batch_size=batch_size
+            net=net, in_shape=in_shape, dataset=dataset, batch_size=batch_size,\
+                min_batches=min_batches, min_seconds=min_seconds
         )

dl_bench/utils.py

@@ -121,6 +121,8 @@ def str_to_dtype(dtype: str):
         return torch.float32
     elif dtype == "bfloat16":
         return torch.bfloat16
+    elif dtype == "int8":
+        return torch.qint8
     else:
         raise ValueError(f"Unsupported data type: {dtype}")
 
@@ -404,9 +406,9 @@ def inference(self, backend: Backend):
             start = time.perf_counter()
             # Duration is inconsistent now
             with tm.timeit("duration_s"):
-                for i, x in enumerate(test_loader):
+                while True:
                     s = get_time()
-                    x = backend.to_device(x)
+                    x = backend.to_device(sample)
                     if backend.dtype != torch.float32:
                         with torch.autocast(
                             device_type=backend.device_name,