Skip to content

Instantly share code, notes, and snippets.

@Termina1
Last active January 30, 2021 12:36
Show Gist options
  • Save Termina1/91f8f1122e3b7516408444b660704e8f to your computer and use it in GitHub Desktop.
Save Termina1/91f8f1122e3b7516408444b660704e8f to your computer and use it in GitHub Desktop.
Multistage docker container for building tensorflow from source Ubuntu 20.04
FROM ubuntu:20.04 as cuda
ENV TZ=Europe/Moscow
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
RUN apt-get update && apt-get install --no-install-recommends -y build-essential git \
wget make build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl \
llvm libncurses5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev ca-certificates gnupg
RUN git clone https://github.com/pyenv/pyenv.git /root/pyenv
RUN /root/pyenv/bin/pyenv install 3.8.0
RUN ln -s /root/.pyenv/versions/3.8.0/bin/python /usr/bin/python
RUN ln -s /root/.pyenv/versions/3.8.0/bin/pip /usr/bin/pip
RUN apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
RUN sh -c 'echo "deb http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64 /" > /etc/apt/sources.list.d/cuda.list'
RUN apt-get update
RUN apt-get install -y cuda-toolkit-11-1
RUN wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64/nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb
RUN apt install ./nvidia-machine-learning-repo-ubuntu2004_1.0.0-1_amd64.deb && apt-get update
RUN apt-get install --no-install-recommends \
libcudnn8=8.0.5.39-1+cuda11.1 \
libcudnn8-dev=8.0.5.39-1+cuda11.1
COPY ./tensorrt /usr/local/tensorrt
FROM cuda as tensorbuild
RUN git clone https://github.com/tensorflow/tensorflow.git
WORKDIR /tensorflow
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/extras/CUPTI/lib64"
RUN pip install -U --user pip "numpy<1.19.0" wheel && pip install -U --user keras_preprocessing --no-deps
RUN curl https://bazel.build/bazel-release.pub.gpg | apt-key add - \
&& echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list \
&& apt update && apt install -y bazel
# RUN perl -pi.bak -e 's%, CompareUFunc%, (PyUFuncGenericFunction) CompareUFunc%g' tensorflow/python/lib/core/bfloat16.cc \
# && sed -i 's/cuda_config.cudart_version/"11.0"/g' third_party/gpus/cuda_configure.bzl
RUN PYTHON_BIN_PATH=$(which python3) \
PYTHON_LIB_PATH="$(python -c 'import site; print(site.getsitepackages()[0])')" \
TF_NEED_OPENCL_SYCL=0 \
TF_NEED_OPENCL=0 \
TF_NEED_ROCM=0 \
TF_NEED_CUDA=1 \
TF_NEED_TENSORRT=1 \
TF_CUDA_VERSION=11 \
TF_CUDNN_VERSION=8 \
TF_TENSORRT_VERSION=7 \
CUDA_TOOLKIT_PATH=/usr/local/cuda \
CUDNN_INSTALL_PATH=/usr/lib/x86_64-linux-gnu \
TENSORRT_INSTALL_PATH=/usr/local/tensorrt \
TF_CUDA_COMPUTE_CAPABILITIES=8.6 \
TF_CUDA_CLANG=0 \
TF_NEED_IGNITE=0 \
TF_ENABLE_XLA=0 \
TF_NEED_MPI=0 \
GCC_HOST_COMPILER_PATH=$(which gcc) \
CC_OPT_FLAGS="-march=native" \
TF_SET_ANDROID_WORKSPACE=0 \
./configure
RUN bazel build --config=cuda //tensorflow/tools/pip_package:build_pip_package
RUN ./bazel-bin/tensorflow/tools/pip_package/build_pip_package --nightly_flag /tmp/tensorflow_pkg
FROM cuda
COPY --from=tensorbuild /tmp/tensorflow_pkg /tmp/tensorflow_pkg
RUN pip install /tmp/tensorflow_pkg/tf_nightly-2.5.0-cp38-cp38-linux_x86_64.whl
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment