Upgrade PyTorch version to v2.2.1 (#2279)

Signed-off-by: Yuki Iwai <[email protected]>
kubeflow · Mar 12, 2024 · 679e6fb · 679e6fb
1 parent 61406a5
commit 679e6fb
Show file tree

Hide file tree

Showing 7 changed files with 9 additions and 11 deletions.
diff --git a/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.cpu b/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.cpu
@@ -6,6 +6,7 @@ ADD examples/v1beta1/trial-images/darts-cnn-cifar10 ${TARGET_DIR}
 
 WORKDIR  ${TARGET_DIR}
 
+RUN pip install --prefer-binary --no-cache-dir torch==2.2.1 torchvision==0.17.1
 RUN pip install --prefer-binary --no-cache-dir -r requirements.txt
 RUN chgrp -R 0 ${TARGET_DIR} \
   && chmod -R g+rwX ${TARGET_DIR}

diff --git a/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu b/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu
@@ -1,7 +1,7 @@
 # We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
-# PyTorch=1.13.0, cuda=11.8.0
-# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-22-11.html#rel-22-11
-FROM nvcr.io/nvidia/pytorch:22.11-py3
+# PyTorch=2.2.0, cuda=12.3.2
+# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-01.html#rel-24-01
+FROM nvcr.io/nvidia/pytorch:24.01-py3
 
 ENV TARGET_DIR /opt/darts-cnn-cifar10
 

diff --git a/examples/v1beta1/trial-images/darts-cnn-cifar10/requirements.txt b/examples/v1beta1/trial-images/darts-cnn-cifar10/requirements.txt
@@ -1,3 +1 @@
-torch==1.13.1
-torchvision==0.14.1
 Pillow>=9.1.1
diff --git a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu
@@ -6,6 +6,7 @@ WORKDIR /opt/pytorch-mnist
 
 # Add folder for the logs.
 RUN mkdir /katib
+RUN pip install --prefer-binary --no-cache-dir torch==2.2.1 torchvision==0.17.1
 RUN pip install --prefer-binary --no-cache-dir -r requirements.txt
 
 RUN chgrp -R 0 /opt/pytorch-mnist \

diff --git a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu
@@ -1,7 +1,7 @@
 # We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
-# PyTorch=1.13.0, cuda=11.8.0
-# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-22-11.html#rel-22-11
-FROM nvcr.io/nvidia/pytorch:22.11-py3
+# PyTorch=2.2.0, cuda=12.3.2
+# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-01.html#rel-24-01
+FROM nvcr.io/nvidia/pytorch:24.01-py3
 
 ADD examples/v1beta1/trial-images/pytorch-mnist /opt/pytorch-mnist
 

diff --git a/examples/v1beta1/trial-images/pytorch-mnist/requirements.txt b/examples/v1beta1/trial-images/pytorch-mnist/requirements.txt
@@ -1,4 +1,2 @@
 cloudml-hypertune==0.1.0.dev6
-torch==1.13.1
-torchvision==0.14.1
 Pillow>=9.1.1
diff --git a/sdk/python/v1beta1/kubeflow/katib/constants/constants.py b/sdk/python/v1beta1/kubeflow/katib/constants/constants.py
@@ -52,7 +52,7 @@
 # TODO (andreyvelich): Implement list_base_images function to get each image description.
 BASE_IMAGE_TENSORFLOW = "docker.io/tensorflow/tensorflow:2.13.0"
 BASE_IMAGE_TENSORFLOW_GPU = "docker.io/tensorflow/tensorflow:2.13.0-gpu"
-BASE_IMAGE_PYTORCH = "docker.io/pytorch/pytorch:1.12.1-cuda11.3-cudnn8-runtime"
+BASE_IMAGE_PYTORCH = "docker.io/pytorch/pytorch:2.2.1-cuda12.1-cudnn8-runtime"
 BASE_IMAGE_MXNET = "docker.io/mxnet/python:1.9.1_native_py3"
 
 DEFAULT_DB_MANAGER_ADDRESS = "katib-db-manager.kubeflow:6789"