diff --git a/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu b/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu index 819b8fb77fa..b5d390dc98f 100644 --- a/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu +++ b/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu @@ -1,6 +1,6 @@ -ARG cuda_version=10.0 -ARG cudnn_version=7 -FROM pytorch/pytorch:1.0-cuda${cuda_version}-cudnn${cudnn_version}-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 ENV TARGET_DIR /opt/darts-cnn-cifar10 diff --git a/scripts/v1beta1/build.sh b/scripts/v1beta1/build.sh index c3aa5d5a0d2..7a1fd94e34a 100755 --- a/scripts/v1beta1/build.sh +++ b/scripts/v1beta1/build.sh @@ -120,16 +120,16 @@ 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 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 . - - echo -e "\nBuilding PyTorch CIFAR-10 CNN training container example for DARTS with GPU support...\n" - docker buildx build --platform linux/amd64 -t "${REGISTRY}/darts-cnn-cifar10-gpu:${TAG}" -f examples/${VERSION}/trial-images/darts-cnn-cifar10/Dockerfile.gpu . - echo -e "\nBuilding dynamic learning rate training container example for PBT...\n" docker buildx build --platform linux/amd64 -t "${REGISTRY}/simple-pbt:${TAG}" -f examples/${VERSION}/trial-images/simple-pbt/Dockerfile . fi +echo -e "\nBuilding PyTorch CIFAR-10 CNN training container example for DARTS with CPU support...\n" +docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/darts-cnn-cifar10-cpu:${TAG}" -f examples/${VERSION}/trial-images/darts-cnn-cifar10/Dockerfile.cpu . + +echo -e "\nBuilding PyTorch CIFAR-10 CNN training container example for DARTS with GPU support...\n" +docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/darts-cnn-cifar10-gpu:${TAG}" -f examples/${VERSION}/trial-images/darts-cnn-cifar10/Dockerfile.gpu . + 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 .