Support Ubuntu 20.04 GPU in Docker build (#1869)

pytorch · Sep 22, 2022 · 4f7caa3 · 4f7caa3
1 parent 3c4ee71
commit 4f7caa3
Showing 1 changed file with 27 additions and 0 deletions.
diff --git a/docker/build_image.sh b/docker/build_image.sh
@@ -8,6 +8,7 @@ DOCKER_FILE="Dockerfile"
 BASE_IMAGE="ubuntu:18.04"
 CUSTOM_TAG=false
 CUDA_VERSION=""
+UBUNTU_VERSION="ubuntu:18.04"
 USE_LOCAL_SERVE_FOLDER=false
 BUILD_WITH_IPEX=false
 
@@ -21,6 +22,7 @@ do
           echo "-g, --gpu specify to use gpu"
           echo "-bt, --buildtype specify to created image for codebuild. Possible values: production, dev, codebuild."
           echo "-cv, --cudaversion specify to cuda version to use"
+          echo "-ub, --ubuntu specify ubuntu version. Possible values: ubuntu:18.04, ubuntu 20.04"
           echo "-t, --tag specify tag name for docker image"
           echo "-lf, --use-local-serve-folder specify this option for the benchmark image if the current 'serve' folder should be used during automated benchmarks"
           echo "-ipex, --build-with-ipex specify to build with intel_extension_for_pytorch"
@@ -63,6 +65,7 @@ do
           BUILD_WITH_IPEX=true
           shift
           ;;
+        # With default ubuntu version 18.04
         -cv|--cudaversion)
           CUDA_VERSION="$2"
           if [ $CUDA_VERSION == "cu116" ];
@@ -90,6 +93,30 @@ do
           shift
           shift
           ;;
+        # CUDA 10 is not supported on Ubuntu 20.04
+        -ub|--ubuntu)
+          UBUNTU_VERSION="$2"
+          if [[ $CUDA_VERSION == "cu116"  &&  $UBUNTU_VERSION == "ubuntu20.04" ]];
+          then 
+            BASE_IMAGE="nvidia/cuda:11.6.0-cudnn8-runtime-ubuntu20.04"
+          elif [[ $CUDA_VERSION == "cu113" && $UBUNTU_VERSION == "ubuntu20.04" ]];
+          then 
+            BASE_IMAGE="nvidia/cuda:11.3.0-cudnn8-runtime-ubuntu20.04"
+          elif [[ $CUDA_VERSION == "cu111" && $UBUNTU_VERSION == "ubuntu20.04" ]];
+          then 
+            BASE_IMAGE="nvidia/cuda:11.1.0-cudnn8-runtime-ubuntu20.04"
+          elif [[ $UBUNTU_VERSION == "ubuntu20.04" ]];
+          then
+            echo "Using CPU image"
+            BASE_IMAGE="ubuntu:20.04"
+          else
+            echo "Ubuntu and CUDA version combination is not supported"
+            echo $UBUNTU_VERSION
+            echo $CUDA_VERSION
+            exit 1 
+          fi
+          shift
+          ;;
     esac
 done