infer with different btch sizes

Author: mikel-brostrom

tests/unit/test_exporters_dynamic.py

@@ -11,7 +11,7 @@
 from boxmot.utils.checks import RequirementsChecker
 
 
-def _load_existing_osnet_model_and_input():
+def _load_existing_osnet_model_and_input(batch_size=2):
     candidates = [
         ROOT / "osnet_x0_25_msmt17.pt",
         WEIGHTS / "osnet_x0_25_msmt17.pt",
@@ -22,7 +22,7 @@ def _load_existing_osnet_model_and_input():
 
     backend = ReidAutoBackend(weights=weights, device="cpu", half=False)
     model = backend.model.model.eval()
-    im = torch.randn(2, 3, 256, 128)
+    im = torch.randn(batch_size, 3, 256, 128)
     return model, im
 
 
@@ -42,75 +42,86 @@ def _disable_dep_sync(monkeypatch):
     monkeypatch.setattr(RequirementsChecker, "sync_extra", lambda *args, **kwargs: None)
 
 
-def test_onnx_export_dynamic_uses_dynamic_shapes(monkeypatch, tmp_path):
+@pytest.mark.parametrize("batch_size", [1, 2, 4])
+def test_onnx_export_dynamic_uses_dynamic_shapes(monkeypatch, tmp_path, batch_size):
     _disable_dep_sync(monkeypatch)
     _install_fake_onnx(monkeypatch)
 
     calls = []
 
     def fake_export(model, args, f, **kwargs):
-        calls.append(kwargs)
+        calls.append((args, kwargs))
         Path(f).touch()
 
     monkeypatch.setattr(torch.onnx, "export", fake_export)
 
-    model, im = _load_existing_osnet_model_and_input()
+    model, im = _load_existing_osnet_model_and_input(batch_size=batch_size)
     out_file = tmp_path / "osnet_x0_25_msmt17.pt"
 
     exporter = ONNXExporter(model, im, out_file, opset=17, dynamic=True, half=False, simplify=False)
     exported = exporter.export()
 
     assert exported == out_file.with_suffix(".onnx")
     assert len(calls) == 1
-    assert "dynamic_shapes" in calls[0]
-    assert "dynamic_axes" not in calls[0]
+    export_args, export_kwargs = calls[0]
+    assert export_args[0].shape[0] == batch_size
+    assert "dynamic_shapes" in export_kwargs
+    assert "dynamic_axes" not in export_kwargs
 
 
-def test_onnx_export_dynamic_fallback_uses_dynamic_axes(monkeypatch, tmp_path):
+@pytest.mark.parametrize("batch_size", [1, 3])
+def test_onnx_export_dynamic_fallback_uses_dynamic_axes(monkeypatch, tmp_path, batch_size):
     _disable_dep_sync(monkeypatch)
     _install_fake_onnx(monkeypatch)
 
     calls = []
 
     def fake_export(model, args, f, **kwargs):
-        calls.append(kwargs)
+        calls.append((args, kwargs))
         if len(calls) == 1:
             raise RuntimeError("force dynamic fallback")
         Path(f).touch()
 
     monkeypatch.setattr(torch.onnx, "export", fake_export)
 
-    model, im = _load_existing_osnet_model_and_input()
+    model, im = _load_existing_osnet_model_and_input(batch_size=batch_size)
     out_file = tmp_path / "osnet_x0_25_msmt17.pt"
 
     exporter = ONNXExporter(model, im, out_file, opset=17, dynamic=True, half=False, simplify=False)
     exported = exporter.export()
 
     assert exported == out_file.with_suffix(".onnx")
     assert len(calls) == 2
-    assert "dynamic_shapes" in calls[0]
-    assert "dynamic_axes" in calls[1]
+    first_args, first_kwargs = calls[0]
+    second_args, second_kwargs = calls[1]
+    assert first_args[0].shape[0] == batch_size
+    assert second_args[0].shape[0] == batch_size
+    assert "dynamic_shapes" in first_kwargs
+    assert "dynamic_axes" in second_kwargs
 
 
-def test_onnx_export_static_has_no_dynamic_shapes(monkeypatch, tmp_path):
+@pytest.mark.parametrize("batch_size", [1, 2, 5])
+def test_onnx_export_static_has_no_dynamic_shapes(monkeypatch, tmp_path, batch_size):
     _disable_dep_sync(monkeypatch)
     _install_fake_onnx(monkeypatch)
 
     calls = []
 
     def fake_export(model, args, f, **kwargs):
-        calls.append(kwargs)
+        calls.append((args, kwargs))
         Path(f).touch()
 
     monkeypatch.setattr(torch.onnx, "export", fake_export)
 
-    model, im = _load_existing_osnet_model_and_input()
+    model, im = _load_existing_osnet_model_and_input(batch_size=batch_size)
     out_file = tmp_path / "osnet_x0_25_msmt17.pt"
 
     exporter = ONNXExporter(model, im, out_file, opset=17, dynamic=False, half=False, simplify=False)
     exported = exporter.export()
 
     assert exported == out_file.with_suffix(".onnx")
     assert len(calls) == 1
-    assert "dynamic_shapes" not in calls[0]
-    assert "dynamic_axes" not in calls[0]
+    export_args, export_kwargs = calls[0]
+    assert export_args[0].shape[0] == batch_size
+    assert "dynamic_shapes" not in export_kwargs
+    assert "dynamic_axes" not in export_kwargs