Update test_template_benchmark.py

sdpython · Feb 13, 2020 · f19a806 · f19a806
1 parent 747dd97
commit f19a806
Showing 1 changed file with 10 additions and 11 deletions.
diff --git a/_unittests/ut_asv_benchmark/test_template_benchmark.py b/_unittests/ut_asv_benchmark/test_template_benchmark.py
@@ -3,14 +3,13 @@
 """
 import os
 import unittest
-from onnx.defs import onnx_opset_version
 try:
     from sklearn.utils._testing import ignore_warnings
 except ImportError:
     from sklearn.utils.testing import ignore_warnings
 from skl2onnx.common.exceptions import MissingShapeCalculator
 from pyquickhelper.pycode import ExtTestCase
-
+from mlprodict.tools.asv_options_helper import benchmark_version
 from mlprodict.asv_benchmark.template.skl_model_classifier import (
     TemplateBenchmarkClassifier
 )
@@ -51,7 +50,7 @@ def test_template_benchmark_classifier(self):
         cl.setup_cache()
         N = 60
         nf = cl.params[2][1]
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         optim = None
         for runtime in ['skl', 'pyrt', 'ort']:
@@ -86,7 +85,7 @@ def test_template_benchmark_classifier_raw_score(self):
         cl.setup_cache()
         N = 60
         nf = cl.params[2][1]
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         optim = None
         for runtime in ['skl', 'pyrt', 'ort']:
@@ -121,7 +120,7 @@ def test_template_benchmark_clustering(self):
         cl.setup_cache()
         N = 60
         nf = cl.params[2][1]
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         optim = None
         for runtime in ['skl', 'pyrt']:
@@ -153,7 +152,7 @@ def test_template_benchmark_regressor(self):
         cl.setup_cache()
         N = 60
         nf = cl.params[2][1]
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         optim = None
         for runtime in ['skl', 'pyrt', 'ort']:
@@ -188,7 +187,7 @@ def test_template_benchmark_multi_classifier(self):
         cl.setup_cache()
         N = 60
         nf = cl.params[2][1]
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         optim = None
         for runtime in ['skl', 'pyrt']:
@@ -225,7 +224,7 @@ def test_template_benchmark_outlier(self):
         N = 60
         nf = cl.params[2][1]
         expect = 10
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         optim = None
         for runtime in ['skl', 'pyrt']:
@@ -264,7 +263,7 @@ def test_template_benchmark_trainable_transform(self):
         cl.setup_cache()
         N = 60
         nf = cl.params[2][1]
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         expect = 12
         optim = None
@@ -305,7 +304,7 @@ def test_template_benchmark_transform(self):
         cl.setup_cache()
         N = 60
         nf = cl.params[2][1]
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         expect = 10
         optim = None
@@ -345,7 +344,7 @@ def test_template_benchmark_transformPositive(self):
         cl.setup_cache()
         N = 60
         nf = cl.params[2][1]
-        opset = onnx_opset_version()
+        opset = benchmark_version()[-1]
         dtype = 'float'
         expect = 12
         optim = None