Windows环境下编译TensorFlow 2.4源码记录

前言

最近两天重装了PC系统，时间也比较多，正好试试在Windows系统上编译TensorFlow源码。我想尝试编译最新的TensorFlow 2.4.1，搭配CUDA 11.2和cuDNN 8.1，看看在最新的版本下编译会不会出现问题。安照官方的教程走一遍后，编译过程还是比较顺利，没有遇到什么麻烦的问题，可以说不管是在Linux还是Windows下，使用最新版本的TensorFlow、CUDA和cuDNN都能顺利通过编译。下面就简单记录下编译的过程，免得以后重复板砖。

安装环境依赖

参考官方的在 Windows 环境中从源代码构建教程，在Windows上编译TensorFlow源码需要安装的环境如下：

• VS2019
• CUDA
• cuDNN
• Bazel
• Anaconda3
• MSYS2

安装过程提示

CUDA下载最新的11.2版本，在安装时可选择自定义，仅安装CUDA选项即可；cuDNN对应的版本是8.1.0，下载后把解压的文件复制到CUDA对应的文件夹中；VS2019在安装时勾选使用C++的桌面开发组件；对于TensorFlow 2.4.x的版本，指定的Bazel版本为3.1.0。Bazel可检测PC上安装的Visual Studio，当安装多个VS版本时，可通过设置BAZEL_VC环境变量使用具体的版本，如要使用VS2019，则设置BAZEL_VC环境变量为C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC。安装好Bazel后，重命名Bazel可执行文件为bazel.exe，添加所在目录到PATH环境变量；安装好Anaconda3后，conda和pip需要切换到国内下载源，以切换到中科大源为例，打开Anaconda命令行，依次输入如下命令更换conda下载源：

conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/
conda config --remove channels defaults
conda config --set show_channel_urls yes

在Windows用户目录下新建名为pip的文件夹，在文件夹内新建文件pip.ini（Linux环境下新建文件夹.pip，在文件夹内新建文件pip.conf），打开文件键入如下内容更换pip下载源：

[global]
index-url = https://pypi.mirrors.ustc.edu.cn/simple/
[install]
trusted-host = https://pypi.mirrors.ustc.edu.cn

创建好Python虚拟环境后，切换到虚拟环境，输入如下命令安装编译TensorFlow需要的包：

pip install six numpy wheel
pip install keras_applications==1.0.6 --no-deps
pip install keras_preprocessing==1.0.5 --no-deps

在安装MSYS2后，可先修改MSYS2的下载源，打开安装目录下的\etc\pacman.d目录，在mirrorlist.mingw32文件的开头添加一行：

Server = https://mirrors.tuna.tsinghua.edu.cn/msys2/mingw/i686/

在mirrorlist.mingw64文件的开头添加一行：

Server = https://mirrors.tuna.tsinghua.edu.cn/msys2/mingw/x86_64/

同样在mirrorlist.msys文件的开头添加一行：

Server = https://mirrors.tuna.tsinghua.edu.cn/msys2/msys/$arch/

添加安装目录下的\usr\bin目录到PATH环境变量，打开cmd控制台，输入如下命令安装编译所需的命令工具：

pacman -Syu
pacman -S git patch unzip zip

编译源码

在安装好环境依赖后，打开Anaconda命令行工具，切换Python虚拟环境下，输入如下命令下载TensorFlow源码：

git clone https://github.com/tensorflow/tensorflow.git

git的方式可能很慢，可离线下载TensorFlow 2.4.1版本的源码。源码下载好后，cd进入根目录，如果是git方式下载的全部版本的源码，则可通过如下命令切换到2.4分支：

git checkout r2.4

接着输入命令python configure.py配置TensorFlow，配置的示例如下：

(base) E:\>cd E:\Downloads\TensorFlow\tensorflow-2.4.1
(base) E:\Downloads\TensorFlow\tensorflow-2.4.1>conda activate tf-nightly
(tf-nightly) E:\Downloads\TensorFlow\tensorflow-2.4.1>python configure.py
You have bazel 3.1.0 installed.
Please specify the location of python. [Default is D:\Anaconda3\envs\tf-nightly\python.exe]:

Found possible Python library paths:
  D:\Anaconda3\envs\tf-nightly\lib\site-packages
Please input the desired Python library path to use.  Default is [D:\Anaconda3\envs\tf-nightly\lib\site-packages]

Do you wish to build TensorFlow with ROCm support? [y/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.

Found CUDA 11.2 in:
    C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.2/lib/x64
    C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.2/include
Found cuDNN 8 in:
    C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.2/lib/x64
    C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.2/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]: 7.5

Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is /arch:AVX]:

Would you like to override eigen strong inline for some C++ compilation to reduce the compilation time? [Y/n]:
Eigen strong inline overridden.

Would you like to interactively configure ./WORKSPACE for Android builds? [y/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=mkl_aarch64    # Build with oneDNN support for Aarch64.
        --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.

(tf-nightly) E:\Downloads\TensorFlow\tensorflow-2.4.1>

配置好后，输入如下命令编译pip软件包：

bazel build --config=opt --config=cuda --define=no_tensorflow_py_deps=true //tensorflow/tools/pip_package:build_pip_package

首次编译大概会花1个小时，编译完成后，执行如下命令生成pip软件包：

bazel-bin\tensorflow\tools\pip_package\build_pip_package ./tensorflow_pkg

执行完成后会在源码根目录下的tensorflow_pkg文件夹下生成.whl文件，如tensorflow-2.4.1-cp36-cp36m-win_amd64.whl。如果需要构建TensorFlow C++运行库，可输入如下命令：

bazel build --config=opt --config=cuda //tensorflow/tools/lib_package:libtensorflow

编译完成后可在\bazel-bin\tensorflow\tools\lib_package目录下找到生成的libtensorflow.tar.gz压缩文件，里面包含相关的库文件。

结语

整个过程比想象的要简单，记得我在刚开始学TensorFlow的时候尝试过在Windows上编译源码，但由于安装的是VS2015，导致编译失败，之后也就懒得再去尝试。目前TensorFlow没有支持到CUDA11.2，要在PC上使用就需要从源码构建TensorFlow，好在这些开源项目有足够好的兼容性，在编译时非常顺利，除了要吐槽Bazel编译时下载包的速度太慢，而且会因为网络问题导致失败，需要重复执行直到相关的软件包都下载完毕。