update batch size for mmdet models

supervisely-ecosystem · Oct 10, 2024 · 3e3cefe · 3e3cefe
1 parent 43dd103
commit 3e3cefe
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/main.py b/src/main.py
@@ -61,7 +61,12 @@ def main_func():
     bm.upload_eval_results(res_dir + "/evaluation/")
 
     try:
-        bm.run_speedtest(g.session_id, g.project_id)
+        batch_sizes = (1, 8, 16)
+        session_info = g.session.get_session_info()
+        support_batch_inference = session_info.get("batch_inference_support", False)
+        if not support_batch_inference:
+            batch_sizes = (1,)
+        bm.run_speedtest(g.session_id, g.project_id, batch_sizes=batch_sizes)
         sec_pbar.hide()
         bm.upload_speedtest_results(res_dir + "/speedtest/")
     except Exception as e: