test

Author: andrey-churkin

tests/post_training/pipelines/image_classification_timm.py

@@ -11,6 +11,7 @@
 
 import numpy as np
 import onnx
+import onnxoptimizer
 import openvino as ov
 import timm
 import torch
@@ -45,17 +46,22 @@ def prepare_model(self) -> None:
             onnx_path = self.fp32_model_dir / "model_fp32.onnx"
             additional_kwargs = {}
             if self.batch_size > 1:
-                additional_kwargs["input_names"] = ["image"]
-                additional_kwargs["dynamic_axes"] = {"image": {0: "batch"}}
+                batch = torch.export.Dim("batch")
+                additional_kwargs["dynamic_shapes"] = ({0: batch},)
+
             torch.onnx.export(
                 timm_model,
                 self.dummy_tensor,
                 onnx_path,
                 export_params=True,
                 opset_version=13,
+                dynamo=True,
                 **additional_kwargs,
             )
-            self.model = onnx.load(onnx_path)
+
+            model = onnx.load(onnx_path)
+            passes = ["fuse_bn_into_conv"]
+            self.model = onnxoptimizer.optimize(model, passes)
             self.input_name = self.model.graph.input[0].name
 
         if self.backend in OV_BACKENDS + [BackendType.FP32]:

tests/post_training/requirements.txt

@@ -23,3 +23,4 @@ transformers==4.53.0
 whowhatbench @ git+https://github.com/openvinotoolkit/[email protected]#subdirectory=tools/who_what_benchmark
 datasets==3.6.0
 onnxscript==0.5.4
+onnxoptimizer==0.3.8