Fix make_lmdb when specify models & topk harness

Change-Id: Id868179fbde38ccab4a0cfe779e9def0b2234b01
Jeffwhen · Aug 2, 2022 · f37ad95 · f37ad95
1 parent 26f2fcb
commit f37ad95
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/python/tpu_perf/harness/topk.py b/python/tpu_perf/harness/topk.py
@@ -50,15 +50,18 @@ def enqueue():
         self.model.put()
 
     def process(self):
+        arg_results = dict()
         while True:
             task_id, results, valid = self.model.get()
             if task_id == 0:
                 break
             output = results[0]
             output = output.reshape(output.shape[0], -1)
-            labels = self.label.pop(task_id)
             argmaxs = np.argmax(output, axis=-1)
             topks = np.argpartition(output, -5)[:, -5:]
+            arg_results[task_id] = (argmaxs, topks)
+        for task_id, (argmaxs, topks) in arg_results.items():
+            labels = self.label.pop(task_id)
             for label, argmax, topk in zip(labels, argmaxs, topks):
                 self.stats['count'] += 1
                 if label == argmax:

diff --git a/python/tpu_perf/make_lmdb.py b/python/tpu_perf/make_lmdb.py
@@ -82,9 +82,9 @@ def main():
                 futures.append(f)
         else:
             for name in sys.argv[1:]:
-                path, config = tree.read_dir(name)
-                f = executor.submit(build_lmdb, data_dir, tree, path, config)
-                futures.append(f)
+                for path, config in tree.read_dir(name):
+                    f = executor.submit(build_lmdb, data_dir, tree, path, config)
+                    futures.append(f)
 
         for f in as_completed(futures):
             err = f.exception()