From 7055eb4ac3ba258fa9e70c5bec283ac476b82cf7 Mon Sep 17 00:00:00 2001 From: tenzen-y Date: Mon, 29 Aug 2022 19:38:24 +0900 Subject: [PATCH] support for arm64 in pytorch-mnist image --- .../trial-images/pytorch-mnist/Dockerfile.gpu | 4 +++- scripts/v1beta1/build.sh | 12 ++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu index 0cc37f2cd5d..4282f51c8c3 100644 --- a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu +++ b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu @@ -1,4 +1,6 @@ -FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime +# Pytorch=1.11.0, cuda=11.6.0 +# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel_22-08.html#rel_22-08 +FROM nvcr.io/nvidia/pytorch:22.02-py3 ADD examples/v1beta1/trial-images/pytorch-mnist /opt/pytorch-mnist diff --git a/scripts/v1beta1/build.sh b/scripts/v1beta1/build.sh index 0acc611e04a..8618e94ec57 100755 --- a/scripts/v1beta1/build.sh +++ b/scripts/v1beta1/build.sh @@ -120,12 +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 PyTorch mnist training container example with CPU support...\n" - docker buildx build --platform linux/amd64 -t "${REGISTRY}/pytorch-mnist-cpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.cpu . - - echo -e "\nBuilding PyTorch mnist training container example with GPU support...\n" - docker buildx build --platform linux/amd64 -t "${REGISTRY}/pytorch-mnist-gpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.gpu . - 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 . @@ -142,6 +136,12 @@ fi echo -e "\nBuilding mxnet mnist training container example...\n" docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/mxnet-mnist:${TAG}" -f examples/${VERSION}/trial-images/mxnet-mnist/Dockerfile . +echo -e "\nBuilding PyTorch mnist training container example with CPU support...\n" +docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/pytorch-mnist-cpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.cpu . + +echo -e "\nBuilding PyTorch mnist training container example with GPU support...\n" +docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/pytorch-mnist-gpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.gpu . + echo -e "\nBuilding Tensorflow with summaries mnist training container example...\n" docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/tf-mnist-with-summaries:${TAG}" -f examples/${VERSION}/trial-images/tf-mnist-with-summaries/Dockerfile .