Upgrade Tensorflow version (#2079)

Signed-off-by: Yuki Iwai <[email protected]> Signed-off-by: Yuki Iwai <[email protected]>
kubeflow · Jan 11, 2023 · d76f01e · d76f01e
1 parent 45a4744
commit d76f01e
Show file tree

Hide file tree

Showing 7 changed files with 11 additions and 14 deletions.
diff --git a/cmd/metricscollector/v1beta1/tfevent-metricscollector/requirements.txt b/cmd/metricscollector/v1beta1/tfevent-metricscollector/requirements.txt
@@ -2,5 +2,4 @@ psutil==5.9.4
 rfc3339>=6.2
 grpcio==1.41.1
 googleapis-common-protos==1.6.0
-tensorflow==2.9.3; platform_machine=="x86_64"
-tensorflow-aarch64==2.9.1; platform_machine=="aarch64"
+tensorflow==2.11.0
diff --git a/cmd/suggestion/nas/enas/v1beta1/requirements.txt b/cmd/suggestion/nas/enas/v1beta1/requirements.txt
@@ -1,5 +1,4 @@
 grpcio==1.41.1
 googleapis-common-protos==1.6.0
 cython>=0.29.24
-tensorflow==2.9.3; platform_machine=="x86_64"
-tensorflow-aarch64==2.9.1; platform_machine=="aarch64"
+tensorflow==2.11.0
diff --git a/examples/v1beta1/trial-images/enas-cnn-cifar10/Dockerfile.gpu b/examples/v1beta1/trial-images/enas-cnn-cifar10/Dockerfile.gpu
@@ -1,6 +1,7 @@
-# tensorflow-gpu=2.9.1, cuda=11.7.1
-# Ref: https://docs.nvidia.com/deeplearning/frameworks/tensorflow-release-notes/rel_22-08.html#rel_22-08
-FROM nvcr.io/nvidia/tensorflow:22.08-tf2-py3
+# We need to use the nvcr.io/nvidia/tensorflow image as a base image to support both linux/amd64 and linux_arm64 platforms.
+# tensorflow-gpu=2.10.1, cuda=11.8.0
+# Ref: https://docs.nvidia.com/deeplearning/frameworks/tensorflow-release-notes/rel-22-12.html#rel-22-12
+FROM nvcr.io/nvidia/tensorflow:22.12-tf2-py3
 
 ENV TARGET_DIR /opt/enas-cnn-cifar10
 ENV PYTHONPATH ${TARGET_DIR}

diff --git a/examples/v1beta1/trial-images/enas-cnn-cifar10/RunTrial.py b/examples/v1beta1/trial-images/enas-cnn-cifar10/RunTrial.py
@@ -63,7 +63,7 @@
         test_model = constructor.build_model()
         test_model.summary()
         test_model.compile(loss=keras.losses.categorical_crossentropy,
-                           optimizer=keras.optimizers.Adam(learning_rate=1e-3, decay=1e-4),
+                           optimizer=keras.optimizers.Adam(learning_rate=1e-3),
                            metrics=['accuracy'])
 
     print(">>> Model Constructed Successfully\n")

diff --git a/examples/v1beta1/trial-images/enas-cnn-cifar10/requirements.txt b/examples/v1beta1/trial-images/enas-cnn-cifar10/requirements.txt
@@ -1,3 +1,2 @@
 scipy>=1.7.2
-tensorflow==2.9.3; platform_machine=="x86_64"
-tensorflow-aarch64==2.9.1; platform_machine=="aarch64"
+tensorflow==2.11.0
diff --git a/examples/v1beta1/trial-images/tf-mnist-with-summaries/requirements.txt b/examples/v1beta1/trial-images/tf-mnist-with-summaries/requirements.txt
@@ -1,2 +1 @@
-tensorflow==2.9.3; platform_machine=="x86_64"
-tensorflow-aarch64==2.9.1; platform_machine=="aarch64"
+tensorflow==2.11.0
diff --git a/sdk/python/v1beta1/kubeflow/katib/constants/constants.py b/sdk/python/v1beta1/kubeflow/katib/constants/constants.py
@@ -32,8 +32,8 @@
 
 # Supported base images for the Katib Trials.
 # TODO (andreyvelich): Implement list_base_images function to get each image description.
-BASE_IMAGE_TENSORFLOW = "docker.io/tensorflow/tensorflow:2.9.1"
-BASE_IMAGE_TENSORFLOW_GPU = "docker.io/tensorflow/tensorflow:2.9.1-gpu"
+BASE_IMAGE_TENSORFLOW = "docker.io/tensorflow/tensorflow:2.11.0"
+BASE_IMAGE_TENSORFLOW_GPU = "docker.io/tensorflow/tensorflow:2.11.0-gpu"
 BASE_IMAGE_PYTORCH = "docker.io/pytorch/pytorch:1.12.1-cuda11.3-cudnn8-runtime"
 BASE_IMAGE_MXNET = "docker.io/mxnet/python:1.9.1_native_py3"