50 lines
1.8 KiB
Docker
50 lines
1.8 KiB
Docker
ARG ARCH=
|
|
ARG CUDA=11.0
|
|
FROM nvidia/cuda${ARCH:+-$ARCH}:${CUDA}-base-ubuntu${UBUNTU_VERSION} as base
|
|
# ARCH and CUDA are specified again because the FROM directive resets ARGs
|
|
# (but their default value is retained if set previously)
|
|
ARG ARCH
|
|
ARG CUDA
|
|
ARG CUDNN=8.0.4.30-1
|
|
ARG CUDNN_MAJOR_VERSION=8
|
|
ARG LIB_DIR_PREFIX=x86_64
|
|
ARG LIBNVINFER=7.1.3-1
|
|
ARG LIBNVINFER_MAJOR_VERSION=7
|
|
|
|
# Needed for string substitution
|
|
SHELL ["/bin/bash", "-c"]
|
|
# Pick up some TF dependencies
|
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
build-essential \
|
|
cuda-command-line-tools-${CUDA/./-} \
|
|
libcublas-${CUDA/./-} \
|
|
cuda-nvrtc-${CUDA/./-} \
|
|
libcufft-${CUDA/./-} \
|
|
libcurand-${CUDA/./-} \
|
|
libcusolver-${CUDA/./-} \
|
|
libcusparse-${CUDA/./-} \
|
|
curl \
|
|
libcudnn8=${CUDNN}+cuda${CUDA} \
|
|
libfreetype6-dev \
|
|
libhdf5-serial-dev \
|
|
libzmq3-dev \
|
|
pkg-config \
|
|
software-properties-common \
|
|
unzip
|
|
|
|
# Install TensorRT if not building for PowerPC
|
|
RUN [[ "${ARCH}" = "ppc64le" ]] || { apt-get update && \
|
|
apt-get install -y --no-install-recommends libnvinfer${LIBNVINFER_MAJOR_VERSION}=${LIBNVINFER}+cuda${CUDA} \
|
|
libnvinfer-plugin${LIBNVINFER_MAJOR_VERSION}=${LIBNVINFER}+cuda${CUDA} \
|
|
&& apt-get clean \
|
|
&& rm -rf /var/lib/apt/lists/*; }
|
|
|
|
# For CUDA profiling, TensorFlow requires CUPTI.
|
|
ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
|
|
|
|
# Link the libcuda stub to the location where tensorflow is searching for it and reconfigure
|
|
# dynamic linker run-time bindings
|
|
RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 \
|
|
&& echo "/usr/local/cuda/lib64/stubs" > /etc/ld.so.conf.d/z-cuda-stubs.conf \
|
|
&& ldconfig
|