TensorFlow和PyTorch是两个最流行的机器学习框架。两者在 Arm 上的使用都在增加,从像 Raspberry Pi 这样的小型系统到用于服务器和高性能计算 (HPC) 的大型系统。尽管包中已经对 AArch64 有一些支持,但用户可能希望从源代码编译所有内容。原因包括使用特定工具、针对不同的运行时环境以及尝试从底层库改进性能。Arm 继续努力使 ML on Arm 得到很好的支持,并做出优化以实现尽可能高的性能。
我们希望这些 Docker 镜像和创建它们的方法对希望在 AArch64 上使用 TensorFlow 和 PyTorch 的人有所帮助。
包括什么?
构建基于 Ubuntu 18.04 的 Docker 映像的脚本可从GitHub 上的 Arm 工具解决方案存储库获得。
完成的 TensorFlow 和 PyTorch 图像包含:
- AArch64 的 Ubuntu 18.04 用户空间
- Gnu gcc 编译器 9.3.0
- 数学库,包括Arm 优化例程和 OpenBLAS 0.3.9
TensorFlow 映像还包含一个从 CPython 3.7 构建的 Python3 环境,其中包含:
- NumPy 1.17.1
- TensorFlow 1.15.2 或 TensorFlow 2.3.0 和 SciPy 1.4.1(取决于所选的 TensorFlow 版本)
- TensorFlow 基准测试
- MLPerf
PyTorch 映像还包含一个从 CPython 3.8 构建的 Python3 环境,其中包含:
- NumPy 1.19.1
- SciPy 1.5.2
- PyTorch 1.6
- MLPerf
入门
要构建和运行 Docker 映像,请确保使用的机器是 Arm AArch64。
$ uname -m
aarch64
任何由Graviton提供支持的 AWS EC2 实例 均可用于试用镜像,包括 A1、T4g、M6g、C6g 或 R6g。TensorFlow 和 PyTorch 图像在 Graviton2 上构建和运行的速度明显更快。要记住的另一件事是,在 Graviton2 上完成的构建将针对 Neoverse-N1 进行优化,并且不适用于 AWS A1 实例。
安装 Docker
推荐使用Docker for Linux。
在 Ubuntu 上为用户名 (ubuntu) 安装 git 和 Docker 的摘要如下:
$ sudo apt update
$ sudo apt upgrade -y
$ curl -fsSL get.docker.com -o get-docker.sh && sh get-docker.sh
$ sudo usermod -aG docker ubuntu ; newgrp docker
$ docker run hello-world
类似的步骤可用于其他 Linux 发行版。
构建图像
首先克隆存储库。如果需要, 安装git 。
$ git clone https://github.com/ARM-software/Tool-Solutions.git
要使用 TensorFlow,请更改到 tensorflow-aarch64/ 目录:
$ cd Tool-Solutions/docker/tensorflow-aarch64
对于 PyTorch,将目录更改为 pytorch-aarch64 目录:
$ cd Tool-Solutions/docker/pytorch-aarch64
每个框架都有一个五阶段的 Dockerfile,因此可以根据需要保存和重用增量进度。
build.sh 脚本构建图像并有一个帮助标志来查看选项。build-type 标志用于指定要构建的特定图像集。
查看脚本和目录以查看构建步骤的详细信息。这些可以根据需要进行修改。
要构建所有图像,请使用:
$ ./build.sh --build-type full
TensorFlow 的默认值为 TensorFlow 1。要构建 TensorFlow 2,请使用命令行选项 --tf_version 2。根据所选的 TensorFlow 版本,图像使用 -v1 或 -v2 标记。
TensorFlow 可以选择使用 oneDNN 构建,使用“--onednn”或“--dnnl”标志。如果没有 '--onednn' 标志,则选择 Tensorflow 的默认 Eigen 后端。oneDNN 的 BLAS 后端也可以使用“--onednn”或“--dnnl”标志来选择。对 C++ 参考内核使用 -onednn reference',对 OpenBLAS 使用'--onednn openblas'。
构建 TensorFlow 容易耗尽内存,但可以使用 bazel_memory_limit 标志来避免耗尽可用内存。
例如,要在具有 32GB 内存的机器上成功构建 TensorFlow 2,请设置以下限制:
$ ./build.sh --build-type full --tf_version 2 --onednn --bazel_memory_limit 30000 --jobs 16
构建图像后,使用 docker 标签和 push 命令将它们保存在您喜欢的图像存储库中。我使用Docker Hub来保存图像。标记图像时替换您的 Docker ID。
$ docker tag tensorflow-v2armpl jasonrandrews/tensorflow-v2armpl
$ docker login (login with username and password)
$ docker push jasonrandrews/tensorflow-v2armpl
接下来,让我们看看如何运行图像。
运行图像
MLCommons (MLPerf)包含在 TensorFlow 和 PyTorch 图像中。
在任何安装了 docker 的 AArch64 机器上,使用以下命令运行基准测试。
$ docker pull jasonrandrews/tensorflow-v2armpl
$ docker tag jasonrandrews/tensorflow-v2armpl tensorflow-v2armpl
$ docker run -it --init tensorflow-v2armpl
示例脚本在最终映像的 $HOME 目录中提供。
要使用 TensorFlow 在 ImageNet 最小验证数据集上运行 resnet50 以进行图像分类:
export DATA_DIR={HOME}/CK-TOOLS/dataset-imagenet-ilsvrc2012-val-min
出口 MODEL_DIR=(密码)
$ ./download-model.sh
$ ./download-dataset.sh
...输出省略
在提示符下选择选项 1:
- imagenet-2012-val-min-resized 版本 2012 (476be5741f52384f)
- imagenet-2012-val-min 版本 2012 (60cb8b2218445c36)
- imagenet-2012-val 版本 2012 (14db79a136d98dd4)
- dataset-imagenet-preprocessed-using-tensorflow (fac1d0d5f4e69a85)
- dataset-imagenet-preprocessed-using-pillow (a6a4613ba6dfd570)
- dataset-imagenet-preprocessed-using-opencv (4932bbdd2ac7a17b)
请选择要安装的包[回车为“0”]:1
...更多输出省略
$ cd ./inference/vision/classification_and_detection
$ ./run_local.sh tf resnet50 cpu
要使用 PyTorch 在 coco 数据集上运行 resnet34 以进行对象检测,请拉取图像并运行。
$ 搬运工拉 jasonrandrews/pytorch
$ 码头工人标签 jasonrandrews/pytorch pytorch
$ docker run -it --init pytorch
export DATA_DIR={HOME}/CK-TOOLS/dataset-coco-2017-val
出口 MODEL_DIR=(密码)
$ ./download-model.sh
$ ./download-dataset.sh
$ cd ./inference/vision/classification_and_detection
$ ./run_local.sh pytorch ssd-resnet34 cpu
设置环境变量 MKLDNN_VERBOSE=1 以验证构建在运行基准测试时使用 oneDNN。
概括
AArch64 上的TensorFlow和PyTorch的Docker 映像 现在可以在 Docker Hub 上获取并快速运行。这些图像使用不同的标签来捕获前面描述的使用各种库的构建选项。审查的说明还使用户能够为这些机器学习框架构建自定义图像。我们欢迎任何反馈,以使它们更易于使用或提高性能。
原作者:杰森·安德鲁斯