编译带有 Intel MKL 的 TensorFlow

我永远也不知道该用哪个版本的 Bazel 编译 TensorFlow;以及……牙膏厂能不能好好写文档。

本文均在 Python3 环境下进行编译,0202 年了为什么还有人在用 Python2。

环境准备

Python 开发环境:

sudo apt install python3-dev python3-pip
pip3 install -U --user pip six numpy wheel setuptools mock 
pip3 install -U --user keras_applications --no-deps
pip3 install -U --user keras_preprocessing --no-deps

Bazel:

这里不直接安装,而是使用 Bazelisk 这个 wrapper 完成对 bazel 的调用。Bazelisk 的 release 为一个 standalone 的二进制包,我们可以直接把它改名为 bazel 放入 PATH 中直接进行调用,所有的参数会自动 passthrough 到具体的 bazel 上。在调用 bazelisk 之前,可以使用 USE_BAZEL_VERSION 这个环境变量指定需要的 bazel 版本,如果没有在本地缓存过,则会被自动下载。由于编译 TensorFlow 对 bazel 版本有严格的要求,使用 bazelisk 可以不用反复重新安装 bazel 就能尝试不同的版本。

Intel MKL:

Intel 有 apt 源可以直接安装这个包

wget https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB
apt-key add GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB
rm GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB

sudo wget https://apt.repos.intel.com/setup/intelproducts.list -O /etc/apt/sources.list.d/intelproducts.list
sudo apt-get update

之后可以在这个页面取得想要安装的包的版本的名字,或者直接用 apt-cache search 也可以。目前最新的 MKL 版本为 intel-mkl-64bit-2020.0-088。

编译

git clone https://github.com/tensorflow/tensorflow.git
cd tensorflow
git checkout branch_name  # r1.9, r1.10, etc.
./configure

在 configure 中完成一些配置,比如是否需要 GPU 支持,MPI 支持等等,这里不赘述。

如果我们要编译一个带有 CPU 优化、MKL 库以及 CUDA 支持的 TensorFlow:

bazel build --config=opt --config=mkl --config=cuda //tensorflow/tools/pip_package:build_pip_package

之后创建 Python 的 .whl 安装包

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

然后在 /tmp/tensorflow_pkg 下就可以看到编译好的安装包了。

Reference

https://www.tensorflow.org/install/source
https://github.com/bazelbuild/bazelisk
https://github.com/tensorflow/tensorflow/issues/26249
https://gist.github.com/pachamaltese/afc4faef2f191b533556f261a46b3aa8
https://software.intel.com/en-us/articles/installing-intel-free-libs-and-python-apt-repo

发表评论

电子邮件地址不会被公开。