libtorch: add option to build singular images (pytorch#643)

Author: seemethere

libtorch/build_all_docker.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+set -eou pipefail
+
+TOPDIR=$(git rev-parse --show-toplevel)
+
+for CUDA_VERSION in 11.2 11.1 11.0 10.2 10.1 cpu; do
+  CUDA_VERSION="${CUDA_VERSION}" conda/build_docker.sh
+done

libtorch/build_docker.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+
+export DOCKER_BUILDKIT=1
+TOPDIR=$(git rev-parse --show-toplevel)
+
+CUDA_VERSION=${CUDA_VERSION:-10.2}
+
+case ${CUDA_VERSION} in
+  cpu)
+    BASE_TARGET=base
+    DOCKER_TAG=cpu
+    ;;
+  *)
+    BASE_TARGET=cuda${CUDA_VERSION}
+    DOCKER_TAG=cuda${CUDA_VERSION}
+    ;;
+esac
+
+(
+  set -x
+  docker build \
+    --target final \
+    --build-arg "BASE_TARGET=${BASE_TARGET}" \
+    --build-arg "CUDA_VERSION=${CUDA_VERSION}" \
+    -t "pytorch/libtorch-cxx11-builder:${DOCKER_TAG}" \
+    -f "${TOPDIR}/libtorch/ubuntu16.04/Dockerfile" \
+    ${TOPDIR}
+)

libtorch/ubuntu16.04/Dockerfile

@@ -1,3 +1,4 @@
+ARG BASE_TARGET=base
 FROM nvidia/cuda:9.2-devel-ubuntu16.04 as base
 
 RUN apt-get clean && apt-get update
@@ -17,48 +18,42 @@ RUN bash ./install_mkl.sh && rm install_mkl.sh
 FROM base as cuda
 ADD ./common/install_cuda.sh install_cuda.sh
 ADD ./common/install_magma.sh install_magma.sh
+ENV CUDA_HOME /usr/local/cuda
 
-FROM cuda as cuda92
+FROM base as conda
+ADD ./common/install_conda.sh install_conda.sh
+RUN bash ./install_conda.sh && rm install_conda.sh
+RUN /opt/conda/bin/conda install -y cmake=3.14
+
+FROM cuda as cuda9.2
 RUN bash ./install_cuda.sh 9.2
 RUN bash ./install_magma.sh 9.2
 
-FROM cuda as cuda101
+FROM cuda as cuda10.1
 RUN bash ./install_cuda.sh 10.1
 RUN bash ./install_magma.sh 10.1
 
-FROM cuda as cuda102
+FROM cuda as cuda10.2
 RUN bash ./install_cuda.sh 10.2
 RUN bash ./install_magma.sh 10.2
 
-FROM cuda as cuda110
+FROM cuda as cuda11.0
 RUN bash ./install_cuda.sh 11.0
 RUN bash ./install_magma.sh 11.0
 
-FROM cuda as cuda111
+FROM cuda as cuda11.1
 RUN bash ./install_cuda.sh 11.1
 RUN bash ./install_magma.sh 11.1
 
-FROM cuda as cuda112
+FROM cuda as cuda11.2
 RUN bash ./install_cuda.sh 11.2
 RUN bash ./install_magma.sh 11.2
 
-FROM base as final
+FROM ${BASE_TARGET} as final
 # Install patchelf
 ADD ./common/install_patchelf.sh install_patchelf.sh
 RUN bash ./install_patchelf.sh && rm install_patchelf.sh
-
+COPY --from=intel /opt/intel /opt/intel
 # Install Anaconda
+COPY --from=conda /opt/conda /opt/conda
 ENV PATH /opt/conda/bin:$PATH
-ADD ./common/install_conda.sh install_conda.sh
-RUN bash ./install_conda.sh && rm install_conda.sh
-RUN /opt/conda/bin/conda install -y cmake=3.14
-
-RUN rm -rf /usr/local/cuda-*
-COPY --from=cuda92   /usr/local/cuda-9.2  /usr/local/cuda-9.2
-COPY --from=cuda101  /usr/local/cuda-10.1 /usr/local/cuda-10.1
-COPY --from=cuda102  /usr/local/cuda-10.2 /usr/local/cuda-10.2
-COPY --from=cuda110  /usr/local/cuda-11.0 /usr/local/cuda-11.0
-COPY --from=cuda111  /usr/local/cuda-11.1 /usr/local/cuda-11.1
-COPY --from=cuda112  /usr/local/cuda-11.2 /usr/local/cuda-11.2
-COPY --from=intel /opt/intel /opt/intel
-ENV CUDA_HOME /usr/local/cuda