Move onnxruntime pin to torch_common

Signed-off-by: Michael Tuttle <quic_mtuttle@quicinc.com>
quic · Mar 8, 2023 · 0824ecb · 0824ecb
1 parent c1ea125
commit 0824ecb
Show file tree

Hide file tree

Showing 6 changed files with 5 additions and 6 deletions.
diff --git a/TrainingExtensions/onnx/test/python/test_qc_quantize_op.py b/TrainingExtensions/onnx/test/python/test_qc_quantize_op.py
@@ -327,7 +327,7 @@ def test_one_shot_quantize_dequantize_cpu_vs_gpu(self):
 
         quant_info_gpu = libquant_info.QcQuantizeInfo()
         quant_node_gpu = helper.make_node(op_name, inputs=['input'], outputs=['output'],
-                                      domain="aimet.customop.cuda", quant_info=libpymo.PtrToInt64(quant_info_gpu))
+                                      domain=op_domain, quant_info=libpymo.PtrToInt64(quant_info_gpu))
         model_gpu = create_model_from_node(quant_node_gpu, input_arr.shape)
         session_gpu = build_session(model_gpu, available_providers)
         qc_op_gpu = QcQuantizeOp(quant_info=quant_info_gpu,

diff --git a/TrainingExtensions/onnx/test/python/test_quantsim.py b/TrainingExtensions/onnx/test/python/test_quantsim.py
@@ -40,6 +40,7 @@
 import torch
 import numpy as np
 from onnx import load_model
+import onnxruntime as ort
 import pytest
 from aimet_common.defs import QuantScheme
 from aimet_onnx.quantsim import QuantizationSimModel
@@ -310,8 +311,8 @@ def onnx_callback(session, inputs):
 
         for node in onnx_sim_gpu.model.graph().node:
             if node.op_type == "QcQuantizeOp":
-                # Note: this check will fail if onnxruntime-gpu is not correctly installed
-                assert node.domain == "aimet.customop.cuda"
+                if 'CUDAExecutionProvider' in ort.get_available_providers():
+                    assert node.domain == "aimet.customop.cuda"
         for node in onnx_sim_cpu.model.graph().node:
             if node.op_type == "QcQuantizeOp":
                 assert node.domain == "aimet.customop.cpu"

diff --git a/packaging/dependencies/reqs_pip_torch_common.txt b/packaging/dependencies/reqs_pip_torch_common.txt
@@ -2,6 +2,7 @@ holoviews==1.12.7
 matplotlib>=3
 numpy<1.24,>=1.16.6
 onnx==1.10.0
+onnxruntime==1.10.0
 onnxruntime-extensions
 onnxsim
 scikit-learn==1.1.3

diff --git a/packaging/dependencies/tf-torch-cpu/reqs_pip_torch_cpu.txt b/packaging/dependencies/tf-torch-cpu/reqs_pip_torch_cpu.txt
@@ -1,5 +1,4 @@
 cumm==0.2.8
-onnxruntime==1.10.0
 spconv==2.1.20
 torch==1.9.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
 torchvision==0.10.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
diff --git a/packaging/dependencies/torch-cpu/reqs_pip_torch_cpu.txt b/packaging/dependencies/torch-cpu/reqs_pip_torch_cpu.txt
@@ -1,5 +1,4 @@
 cumm==0.2.8
-onnxruntime==1.10.0
 spconv==2.1.20
 torch==1.9.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
 torchvision==0.10.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
diff --git a/packaging/dependencies/torch-gpu/reqs_pip_torch_gpu.txt b/packaging/dependencies/torch-gpu/reqs_pip_torch_gpu.txt
@@ -1,5 +1,4 @@
 cumm-cu111==0.2.8
-onnxruntime-gpu==1.10.0
 spconv-cu111==2.1.20
 torch==1.9.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html
 torchvision==0.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html