TensorFlow2.3.1のpipパッケージをDockerコンテナ上でビルドする

TensorFlow 2.3.1はCUDA10.1に対応しており、wheelファイルも公開されていますが、なんらかの理由でCUDA10.2で環境構築してしまっていた場合は自分でビルドする必要があります。ここではnvidia-docker2をインストールしている前提で、Python3.6のgpu版pipパッケージをビルドする手順について説明します。

Dockerfile

FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04

RUN apt-get update
RUN apt-get install -y --no-install-recommends \
    git
RUN curl -OL https://github.com/bazelbuild/bazelisk/releases/download/v1.7.4/bazelisk-linux-amd64 \
    && mv bazelisk-linux-amd64  /usr/local/bin/bazel && chmod +x /usr/local/bin/bazel
WORKDIR /build
RUN git clone https://github.com/tensorflow/tensorflow.git
RUN cd tensorflow && git checkout -b v2.3.1 refs/tags/v2.3.1
COPY requirements.txt /build/tensorflow
RUN apt-get install -y --no-install-recommends python3-pip
RUN pip3 install --upgrade pip
RUN pip3 install setuptools
RUN apt-get install -y --no-install-recommends python3-dev
RUN ln -s /usr/bin/python3.6 /usr/bin/python
RUN pip3 install -r /build/tensorflow/requirements.txt

requirements.txtはsetup.pyのREQUIRED_PACKAGESから転記したものです(これらをインストールする必要があることは、公式の「Build from source」に書かれている通りです。日本語版は情報が古いのと、ドキュメントにリンクされているsetup.pyはmasterブランチなので、ビルドするタグ上に切り替える必要があるので、注意が必要です)。

requirements.txt

absl-py >= 0.7.0
astunparse == 1.6.3
gast == 0.3.3
google_pasta >= 0.1.8
h5py >= 2.10.0, < 2.11.0
keras_preprocessing >= 1.1.1, < 1.2
numpy >= 1.16.0, < 1.19.0
opt_einsum >= 2.3.2
protobuf >= 3.9.2
tensorboard >= 2.3.0, < 3
tensorflow_estimator >= 2.3.0, < 2.4.0
termcolor >= 1.1.0
wrapt >= 1.11.1
wheel >= 0.26
six >= 1.12.0

このDockerfileからimageをビルドし、コンテナを起動します。

docker build -t tfbuild .
docker run --gpus 0 --rm -it  --shm-size=2g  -u root -v "$(pwd):/work" tfbuild bash -l

dockerコンテナのbashが立ち上がったら、次のようにしてビルドします。ただし--local_ram_resources、--local_cpu_resourcesは環境に合わせてください。PCの性能のぎりぎりを指定するとリソースを使い果たしてフリーズすることが多いです。うっかりChromeなど、他のアプリを立ち上げたままビルドしてもフリーズするので、気をつけたほうが良いと思います。

cd tensorflow/
./configure 
bazel build --local_ram_resources=1024 --local_cpu_resources=8 --verbose_failures --config=opt --config=cuda //tensorflow/tools/pip_package:build_pip_package
./bazel-bin/tensorflow/tools/pip_package/build_pip_package /work/pip_package

bazel buildが終わるとwheelファイルを生成するのに必要なbuild_pip_packageが生成されます。wheelファイルはうっかり消さないように、ホストのディレクトリにマウントさせている/work上に生成させておきます。

./configureは自分の環境に合わせて答えれば良いのですが、nvidia製のgpuが入ったPCの場合は、次のように回答すれば良いと思います。

$ /build/tensorflow# ./configure 
You have bazel 3.1.0 installed.
Please specify the location of python. [Default is /usr/bin/python3]: 


Found possible Python library paths:
  /usr/local/lib/python3.6/dist-packages
  /usr/lib/python3/dist-packages
Please input the desired Python library path to use.  Default is [/usr/local/lib/python3.6/dist-packages]

Do you wish to build TensorFlow with OpenCL SYCL support? [y/N]: N
No OpenCL SYCL support will be enabled for TensorFlow.

Do you wish to build TensorFlow with ROCm support? [y/N]: N
No ROCm support will be enabled for TensorFlow.

Do you wish to build TensorFlow with CUDA support? [y/N]: y
CUDA support will be enabled for TensorFlow.

Do you wish to build TensorFlow with TensorRT support? [y/N]: N
No TensorRT support will be enabled for TensorFlow.

Found CUDA 10.2 in:
    /usr/local/cuda-10.2/targets/x86_64-linux/lib
    /usr/local/cuda-10.2/targets/x86_64-linux/include
Found cuDNN 7 in:
    /usr/lib/x86_64-linux-gnu
    /usr/include


Please specify a list of comma-separated CUDA compute capabilities you want to build with.
You can find the compute capability of your device at: https://developer.nvidia.com/cuda-gpus. Each capability can be specified as "x.y" or "compute_xy" to include both virtual and binary GPU code, or as "sm_xy" to only include the binary code.
Please note that each additional compute capability significantly increases your build time and binary size, and that TensorFlow only supports compute capabilities >= 3.5 [Default is: 3.5,7.0]: 


Do you want to use clang as CUDA compiler? [y/N]: N
nvcc will be used as CUDA compiler.

Please specify which gcc should be used by nvcc as the host compiler. [Default is /usr/bin/gcc]: 


Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native -Wno-sign-compare]: 


Would you like to interactively configure ./WORKSPACE for Android builds? [y/N]: N
Not configuring the WORKSPACE for Android builds.

Preconfigured Bazel build configs. You can use any of the below by adding "--config=<>" to your build command. See .bazelrc for more details.
	--config=mkl         	# Build with MKL support.
	--config=monolithic  	# Config for mostly static monolithic build.
	--config=ngraph      	# Build with Intel nGraph support.
	--config=numa        	# Build with NUMA support.
	--config=dynamic_kernels	# (Experimental) Build kernels into separate shared objects.
	--config=v2          	# Build TensorFlow 2.x instead of 1.x.
Preconfigured Bazel build configs to DISABLE default on features:
	--config=noaws       	# Disable AWS S3 filesystem support.
	--config=nogcp       	# Disable GCP support.
	--config=nohdfs      	# Disable HDFS support.
	--config=nonccl      	# Disable NVIDIA NCCL support.
Configuration finished