Support for arm64 in enas-cnn-cifar10 image (#1944)

Signed-off-by: tenzen-y <yuki.iwai.tz@gmail.com> Signed-off-by: tenzen-y <yuki.iwai.tz@gmail.com>
kubeflow · Aug 31, 2022 · daf5b9b · daf5b9b
1 parent 6870587
commit daf5b9b
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 6 deletions.
diff --git a/examples/v1beta1/trial-images/enas-cnn-cifar10/Dockerfile.cpu b/examples/v1beta1/trial-images/enas-cnn-cifar10/Dockerfile.cpu
@@ -8,7 +8,7 @@ ADD examples/v1beta1/trial-images/enas-cnn-cifar10 ${TARGET_DIR}
 
 WORKDIR  ${TARGET_DIR}
 
-RUN if [ "${TARGETARCH}" = "aarch64" ]; then \
+RUN if [ "${TARGETARCH}" = "arm64" ]; then \
     apt-get -y update && \
     apt-get -y install gfortran libpcre3 libpcre3-dev && \
     apt-get clean && \

diff --git a/examples/v1beta1/trial-images/enas-cnn-cifar10/Dockerfile.gpu b/examples/v1beta1/trial-images/enas-cnn-cifar10/Dockerfile.gpu
@@ -1,4 +1,6 @@
-FROM tensorflow/tensorflow:2.9.1-gpu
+# 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
 
 ENV TARGET_DIR /opt/enas-cnn-cifar10
 ENV PYTHONPATH ${TARGET_DIR}

diff --git a/examples/v1beta1/trial-images/tf-mnist-with-summaries/Dockerfile b/examples/v1beta1/trial-images/tf-mnist-with-summaries/Dockerfile
@@ -6,7 +6,7 @@ ADD examples/v1beta1/trial-images/tf-mnist-with-summaries /opt/tf-mnist-with-sum
 
 WORKDIR /opt/tf-mnist-with-summaries
 
-RUN if [ "${TARGETARCH}" = "aarch64" ]; then \
+RUN if [ "${TARGETARCH}" = "arm64" ]; then \
     apt-get -y update && \
     apt-get -y install gfortran libpcre3 libpcre3-dev && \
     apt-get clean && \

diff --git a/scripts/v1beta1/build.sh b/scripts/v1beta1/build.sh
@@ -120,9 +120,6 @@ echo -e "\nBuilding training container images..."
 if [ ! "$ARCH" = "amd64" ]; then
   echo -e "\nSome training container images are supported only amd64."
 else
-  echo -e "\nBuilding Keras CIFAR-10 CNN training container example for ENAS with GPU support...\n"
-  docker buildx build --platform linux/amd64 -t "${REGISTRY}/enas-cnn-cifar10-gpu:${TAG}" -f examples/${VERSION}/trial-images/enas-cnn-cifar10/Dockerfile.gpu .
-
   echo -e "\nBuilding PyTorch CIFAR-10 CNN training container example for DARTS with CPU support...\n"
   docker buildx build --platform linux/amd64 -t "${REGISTRY}/darts-cnn-cifar10-cpu:${TAG}" -f examples/${VERSION}/trial-images/darts-cnn-cifar10/Dockerfile.cpu .
 
@@ -148,4 +145,7 @@ docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/tf-mnist-with-summa
 echo -e "\nBuilding Keras CIFAR-10 CNN training container example for ENAS with CPU support...\n"
 docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/enas-cnn-cifar10-cpu:${TAG}" -f examples/${VERSION}/trial-images/enas-cnn-cifar10/Dockerfile.cpu .
 
+echo -e "\nBuilding Keras CIFAR-10 CNN training container example for ENAS with GPU support...\n"
+docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/enas-cnn-cifar10-gpu:${TAG}" -f examples/${VERSION}/trial-images/enas-cnn-cifar10/Dockerfile.gpu .
+
 echo -e "\nAll Katib images with ${TAG} tag have been built successfully!\n"