gemini-code-assist[bot] commented on code in PR #18787: URL: https://github.com/apache/tvm/pull/18787#discussion_r2813410792
########## docker/install/ubuntu_install_onnx.sh: ########## @@ -20,25 +20,70 @@ set -e set -u set -o pipefail +# We need to fix the onnx version because changing versions tends to break tests +# TODO(mbrookhart): periodically update + +# 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. + # Get the Python version PYTHON_VERSION=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") # Set default value for first argument DEVICE=${1:-cpu} # Install the onnx package -pip3 install \ - onnx==1.20.1 \ - onnxruntime==1.23.2 \ - onnxoptimizer==0.4.2 +pip3 install future + +if [ "$PYTHON_VERSION" == "3.9" ]; then + pip3 install \ + onnx==1.16.0 \ + onnxruntime==1.19.2 \ + onnxoptimizer==0.2.7 -if [ "$DEVICE" == "cuda" ]; then + if [ "$DEVICE" == "cuda" ]; then + pip3 install \ + torch==2.7.0 \ + torchvision==0.22.0 \ + --index-url https://download.pytorch.org/whl/cu118 + else + pip3 install \ + torch==2.7.0 \ + torchvision==0.22.0 \ + --extra-index-url https://download.pytorch.org/whl/cpu + fi +elif [ "$PYTHON_VERSION" == "3.11" ]; then pip3 install \ - torch==2.10.0 \ - torchvision==0.25.0 + onnx==1.17.0 \ + onnxruntime==1.20.1 \ + onnxoptimizer==0.2.7 + + if [ "$DEVICE" == "cuda" ]; then + pip3 install \ + torch==2.7.0 \ + torchvision==0.22.0 \ + --index-url https://download.pytorch.org/whl/cu118 + else + pip3 install \ + torch==2.7.0 \ + torchvision==0.22.0 \ + --extra-index-url https://download.pytorch.org/whl/cpu + fi else pip3 install \ - torch==2.10.0 \ - torchvision==0.25.0 \ - --extra-index-url https://download.pytorch.org/whl/cpu + onnx==1.12.0 \ + onnxruntime==1.12.1 \ + onnxoptimizer==0.2.7 + + if [ "$DEVICE" == "cuda" ]; then + pip3 install \ + torch==2.4.1 \ + torchvision==0.19.1 + else + pip3 install \ + torch==2.4.1 \ + torchvision==0.19.1 \ + --extra-index-url https://download.pytorch.org/whl/cpu + fi fi Review Comment:  This `if/elif/else` block contains significant code duplication for installing `torch` and `onnx` packages across different Python versions. You can refactor this by defining package version strings in variables based on the Python version, and then using these variables in shared `pip3 install` commands. This will make the script more readable and easier to maintain in the future. ```shell if [ "$PYTHON_VERSION" == "3.9" ]; then ONNX_PKGS="onnx==1.16.0 onnxruntime==1.19.2 onnxoptimizer==0.2.7" elif [ "$PYTHON_VERSION" == "3.11" ]; then ONNX_PKGS="onnx==1.17.0 onnxruntime==1.20.1 onnxoptimizer==0.2.7" else ONNX_PKGS="onnx==1.12.0 onnxruntime==1.12.1 onnxoptimizer==0.2.7" fi pip3 install ${ONNX_PKGS} if [ "$PYTHON_VERSION" == "3.9" ] || [ "$PYTHON_VERSION" == "3.11" ]; then TORCH_PKGS="torch==2.7.0 torchvision==0.22.0" TORCH_CUDA_ARGS="--index-url https://download.pytorch.org/whl/cu118"; else TORCH_PKGS="torch==2.4.1 torchvision==0.19.1" TORCH_CUDA_ARGS="" fi if [ "$DEVICE" == "cuda" ]; then pip3 install ${TORCH_PKGS} ${TORCH_CUDA_ARGS} else pip3 install ${TORCH_PKGS} --extra-index-url https://download.pytorch.org/whl/cpu fi ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
