Upgrading onnx and onnxrt verions (#17655)

* Upgrading onnx and onnxrt verions

Upgrading onnx and onnxrt verions

Signed-off-by: Shubham Laddha <[email protected]>

* Upgrading onnx and onnxrt verions

Upgrading onnx and onnxrt verions

Signed-off-by: Shubham Laddha <[email protected]>

---------

Signed-off-by: Shubham Laddha <[email protected]>
Co-authored-by: Shubham Laddha <[email protected]>

Author: Shubhamladdha

docker/install/ubuntu_install_onnx.sh

@@ -26,16 +26,41 @@ set -o pipefail
 # onnx 1.9 removed onnx optimizer from the main repo (see
 # https://github.com/onnx/onnx/pull/2834).  When updating the CI image
 # to onnx>=1.9, onnxoptimizer should also be installed.
-pip3 install \
-    onnx==1.12.0 \
-    onnxruntime==1.12.1 \
-    onnxoptimizer==0.2.7
 
-# torch depends on a number of other packages, but unhelpfully, does
-# not expose that in the wheel!!!
+# Get the Python version
+PYTHON_VERSION=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
+
+# Install the onnx package
 pip3 install future
 
-pip3 install \
-    torch==2.4.1 \
-    torchvision==0.19.1 \
-    --extra-index-url https://download.pytorch.org/whl/cpu
+if [ "$PYTHON_VERSION" == "3.9" ]; then
+    pip3 install \
+        onnx==1.16.0 \
+        onnxruntime==1.19.2 \
+        onnxoptimizer==0.2.7
+
+    pip3 install \
+        torch==2.6.0 \
+        torchvision==0.21.0 \
+        --extra-index-url https://download.pytorch.org/whl/cpu
+elif [ "$PYTHON_VERSION" == "3.11" ]; then
+    pip3 install \
+        onnx==1.17.0 \
+        onnxruntime==1.20.1 \
+        onnxoptimizer==0.2.7
+
+    pip3 install \
+        torch==2.6.0 \
+        torchvision==0.21.0 \
+        --extra-index-url https://download.pytorch.org/whl/cpu
+else
+    pip3 install \
+        onnx==1.12.0 \
+        onnxruntime==1.12.1 \
+        onnxoptimizer==0.2.7
+
+    pip3 install \
+        torch==2.4.1 \
+        torchvision==0.19.1 \
+        --extra-index-url https://download.pytorch.org/whl/cpu
+fi

docker/install/ubuntu_install_python.sh

@@ -34,9 +34,9 @@ fi
 PYTHON_VERSION=$1
 
 case "$PYTHON_VERSION" in
-    3.7|3.8|3.9) ;;
+    3.7|3.8|3.9|3.10|3.11) ;;
     *)
-        echo "Only 3.7, 3.8, and 3.9 versions are supported in this script."
+        echo "Only 3.7, 3.8, 3.9, 3.10 and 3.11 versions are supported in this script."
         exit -1
         ;;
 esac
@@ -58,7 +58,10 @@ case "${release}" in
         [ "${PYTHON_VERSION}" == "3.7" ] && add-apt-repository -y ppa:deadsnakes/ppa
         ;;
     jammy)
-        if [ "${PYTHON_VERSION}" == "3.8" ] || [ "${PYTHON_VERSION}" == "3.9" ]; then
+        if [ "${PYTHON_VERSION}" == "3.8" ] || \
+           [ "${PYTHON_VERSION}" == "3.9" ] || \
+           [ "${PYTHON_VERSION}" == "3.10" ] || \
+           [ "${PYTHON_VERSION}" == "3.11" ]; then
             add-apt-repository -y ppa:deadsnakes/ppa
         fi
         ;;
@@ -98,17 +101,19 @@ fi
 
 # Update pip to match version used to produce requirements-hashed.txt. This step
 # is necessary so that pip's dependency solver is recent.
-pip_spec=$(tac /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt | grep -m 1 'pip==')
-${TVM_VENV}/bin/pip install -U --require-hashes -r <(echo "${pip_spec}") \
-     -c /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt
-
-# Python configuration
-${TVM_VENV}/bin/pip config set global.no-cache-dir true  # Never cache packages
-
-# Now install the remaining base packages.
-${TVM_VENV}/bin/pip install \
-     --require-hashes \
-     -r /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt
+if [ -f "/install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt" ]; then
+    pip_spec=$(tac /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt | grep -m 1 'pip==')
+    ${TVM_VENV}/bin/pip install -U --require-hashes -r <(echo "${pip_spec}") \
+        -c /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt
+
+    # Python configuration
+    ${TVM_VENV}/bin/pip config set global.no-cache-dir true  # Never cache packages
+
+    # Now install the remaining base packages.
+    ${TVM_VENV}/bin/pip install \
+        --require-hashes \
+        -r /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt
+fi
 
 addgroup tvm-venv
 chgrp -R tvm-venv "${TVM_VENV}"