使用TVM优化PyTorch模型实现快速CPU推理
使用TVM优化PyTorch模型实现快速CPU推理
资源不够压榨来凑。没钱加 GPU？推理太慢？只好想办法把 CPU 榨干啦。
作者：Aleksey Bilogur 编译：McGL
Apache TVM 是一个相对较新的 Apache 项目，以深度学习模型推理的性能大幅改进为目标。它属于一种叫做模型编译器(model compilers) 的新技术: 它以高级框架(如 PyTorch 或 TensorFlow)中编写的模型作为输入，生成一个为在特定硬件平台上运行而优化的二进制包作为输出。
在这篇文章中，我们将通过 TVM 的每个步骤来了解它。我们将讨论它工作原理背后的基本概念，然后安装它并运行测试模型进行基准测试。
文章中的完整代码，可以查看 GitHub repo：https://github.com/spellml/examples/tree/master/external/tvm
基本概念
TVM 的核心是一个模型编译器。
如果你熟悉编译编程语言，那么你已经知道编译编程语言几乎严格地比解释编程语言快(想想 C 和 Python)。这是因为增加的编译器步骤允许优化，包括代码的高级表示(例如，循环展开)和低级执行(例如，强制操作对象与硬件处理器原生支持的类型之间的转换) ，这使得代码的执行速度更快，快了一个数量级。
模型编译的目标非常相似: 使用易于编写的高级框架(比如 PyTorch)编写模型。然后，将它的计算图编译成一个二进制对象，该对象只为在一个特定的硬件平台上运行而优化。TVM 支持非常广泛的不同目标平台 —— 有些是预料中的，有些很特别。TVM 文档中的 Getting Started 页面展示了以下支持的后端的图表:
TVM 支持的平台范围绝对是这个项目的优势。例如，PyTorch 的模型量化 API 只支持两个目标平台: x86和 ARM。而使用 TVM，你可以编译模型原生运行在 macOS、 NVIDIA CUDA 上，甚至可以通过 WASM 运行在网络浏览器上。
生成优化模型二进制文件的过程的开始是，将计算图转换成 TVM 的内部高级图格式 —— Relay。Relay 是一个可用的高级模型 API，你甚至可以在其中从头构建新模型，但它主要作为进一步优化模型的统一起点。
TVM 在 Relay 层对图应用一些高级别的优化，然后通过一个叫做“Relay Fusion Pass”的过程将其降低到一个叫做“Tensor Expressions/TE”的低级别 IR。在 TE 层，计算图被分成一组子图，这些子图被 TVM 引擎确定为很好的优化目标。
TVM 优化过程中最后也是最重要的一步是调优。在调优步骤中，TVM 对图中的计算任务(“调度”)的操作顺序进行预测，以在选定的硬件平台上获得最高性能(最快推理时间)。
有趣的是，这不是一个确定性问题 —— 在任何给定的硬件平台上，任何特定操作的执行速度，都存在太多有效的可能排序和太多的不确定性，而且需要考虑到所有其他正在运行的计算过程。TVM 在计算空间上构造了一个搜索空间，然后在该空间上运行一个带有自定义损失函数的 XGBoost 模型，以找到最佳调度方案。
如果这看起来非常复杂，那是因为它本身复杂。幸运的是，你不必知道 TVM 如何工作的任何细节，因为它的高级 API 为你处理大部分细节。
安装 TVM
为了了解 TVM 的性能优势，我编译了一个在 CIFAR10 上进行训练的简单 PyTorch Mobilenet 模型，并测试了它在 TVM 编译之前和之后的推理时间。本文的其余部分将介绍这些代码。你可以在 GitHub 查看：https://github.com/spellml/mobilenet-cifar10。
但是，在使用 TVM 之前，你必须首先安装它。不幸的是，这根本不是一个简单的过程。TVM 目前没有发布任何 wheels，官方文档介绍的是从源代码一步步安装 TVM。
为了测试的目的，我在 AWS 上使用一个 c5.4xlarge 的 CPU 实例。这是一台 x86 机器，因此我们需要同时安装 TVM 和最新版本的 LLVM 工具链。从源代码编译 TVM 大约需要10分钟，所以这是自定义 Docker 镜像的完美用例，我们可以将 TVM 及其所有依赖项编译成一个 Docker 镜像，然后在以后的所有运行中重复使用该镜像。
下面是我使用的 Dockerfile:
FROM ubuntu:18.04
WORKDIR /spell
# Conda install part
RUN apt-get update && \\
apt-get install -y wget git && rm -rf /var/lib/apt/lists/*
ENV CONDA_HOME=/root/anaconda/
RUN wget \\
<https://repo.anaconda.com/miniconda/Miniconda3-py37_4.8.3-Linux-x86_64.sh> \\
&& mkdir /root/.conda \\
&& bash Miniconda3-py37_4.8.3-Linux-x86_64.sh -fbp $CONDA_HOME \\
&& rm -f Miniconda3-py37_4.8.3-Linux-x86_64.sh
ENV PATH=/root/anaconda/bin:$PATH
# NOTE: Spell runs will fail if pip3 is not avaiable at the command line.
# conda injects pip onto the path, but not pip3, so we create a symlink.
RUN ln /root/anaconda/bin/pip /root/anaconda/bin/pip3
# TVM install part
COPY environment.yml /tmp/environment.yml
RUN conda env create -n spell -f=/tmp/environment.yml
COPY scripts/install_tvm.sh /tmp/install_tvm.sh
RUN chmod +x /tmp/install_tvm.sh && /tmp/install_tvm.sh
这使用以下 conda environment.yml:
name: spell
channels:
- conda-forge
dependencies:
- numpy
- pandas
- tornado
- pip
- pip:
# NOTE(aleksey): because of AskUbuntu#1334667, we need an old version of
# XGBoost, as recent versions are not compatible with our base image,
# Ubuntu 18.04. XGBoost is required in this environment because TVM uses
# it as its search space optimization algorithm in the tuning pass.
- xgboost==1.1.0
- torch==1.8.1
- torchvision
- cloudpickle
- psutil
- spell
- kaggle
- tokenizers
- transformers
# NOTE(aleksey): this dependency on pytest is probably accidental, as
# it isn't documented. But without it, the TVM Python package will not
# import.
- pytest
下面是 install_tvm.sh 的内容。请注意，TVM 构建时间变量设置在 config.cmake 文件中，我在这里修改这个文件是为了指向我们使用 apt-get 安装的特定版本的 LLVM:
#!/bin/bash
set -ex
# <https://tvm.apache.org/docs/install/from_source.html#install-from-source>
if [[ ! -d "/tmp/tvm" ]]; then
git clone --recursive <https://github.com/apache/tvm> /tmp/tvm
fi
apt-get update && \
apt-get install -y gcc libtinfo-dev zlib1g-dev \
build-essential cmake libedit-dev libxml2-dev \
llvm-6.0 \
libgomp1 \
zip unzip
if [[ ! -d "/tmp/tvm/build" ]]; then
mkdir /tmp/tvm/build
fi
cp /tmp/tvm/cmake/config.cmake /tmp/tvm/build
mv /tmp/tvm/build/config.cmake /tmp/tvm/build/~config.cmake && \
cat /tmp/tvm/build/~config.cmake | \
# sed -E "s|set\(USE_CUDA OFF\)|set\(USE_CUDA ON\)|" | \
sed -E "s|set\(USE_GRAPH_RUNTIME OFF\)|set\(USE_GRAPH_RUNTIME ON\)|" | \
sed -E "s|set\(USE_GRAPH_RUNTIME_DEBUG OFF\)|set\(USE_GRAPH_RUNTIME_DEBUG ON\)|" | \
sed -E "s|set\(USE_LLVM OFF\)|set\(USE_LLVM /usr/bin/llvm-config-6.0\)|" > \\
/tmp/tvm/build/config.cmake
cd /tmp/tvm/build && cmake .. && make -j4
cd /tmp/tvm/python && /root/anaconda/envs/spell/bin/python setup.py install --user && cd ..
你可以在本地机器上构建 docker，然后使用这个 Gist(https://gist.github.com/ResidentMario/9f41ac480f9efbf2ff1d05d450c29470) 在 EC2 机器上构建这个映像，或者重用我为这个演示构建的公共镜像(https://hub.docker.com/r/residentmario/tvm)，可以完全跳过这个编译过程。
使用 TVM 编译模型
安装了 TVM 之后，我们可以继续使用它编译测试模型。
请注意，TVM 两种客户端，Python 和 CLI; 我在这个项目中使用了 Python 客户端。
首先，我们需要一个训练好的模型。事实上，并不是任意模型都可以。相关的方法 tvm.relay.frontend.from_pytorch只接受一个量化模型作为输入。
量化是将模型图中的操作降低到较低精度表示(例如从 fp32 降低到 int8)的过程。这是模型性能优化的一种形式: 操作数的比特越少，操作的速度就越快。量化是一项复杂的技术，本身也是比较新的技术，在编写本文时，其 PyTorch 实现(torch.jit 模块)仍处于 beta 阶段。我们以前的文章已经深入讨论过量化，所以这里略过这些细节。
从代码中可以看到量化后的模型定义:
def conv_bn(inp, oup, stride):
return nn.Sequential(OrderedDict([
('q', torch.quantization.QuantStub()),
('conv2d', nn.Conv2d(inp, oup, 3, stride, 1, bias=False)),
('batchnorm2d', nn.BatchNorm2d(oup)),
('relu6', nn.ReLU6(inplace=True)),
('dq', torch.quantization.DeQuantStub())
]))
def conv_1x1_bn(inp, oup):
return nn.Sequential(OrderedDict([
('q', torch.quantization.QuantStub()),
('conv2d', nn.Conv2d(inp, oup, 1, 1, 0, bias=False)),
('batchnorm2d', nn.BatchNorm2d(oup)),
('relu6', nn.ReLU6(inplace=True)),
('dq', torch.quantization.DeQuantStub())
]))
def make_divisible(x, divisible_by=8):
import numpy as np
return int(np.ceil(x * 1. / divisible_by) * divisible_by)
class InvertedResidual(nn.Module):
def __init__(self, inp, oup, stride, expand_ratio):
super(InvertedResidual, self).__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = int(inp * expand_ratio)
self.use_res_connect = self.stride == 1 and inp == oup
if expand_ratio == 1:
self.conv = nn.Sequential(OrderedDict([
('q', torch.quantization.QuantStub()),
# dw
('conv2d_1', nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False)),
('bnorm_2', nn.BatchNorm2d(hidden_dim)),
('relu6_3', nn.ReLU6(inplace=True)),
# pw-linear
('conv2d_4', nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False)),
('bnorm_5', nn.BatchNorm2d(oup)),
('dq', torch.quantization.DeQuantStub())
]))
else:
self.conv = nn.Sequential(OrderedDict([
('q', torch.quantization.QuantStub()),
# pw
('conv2d_1', nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False)),
('bnorm_2', nn.BatchNorm2d(hidden_dim)),
('relu6_3', nn.ReLU6(inplace=True)),
# dw
('conv2d_4', nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False)),
('bnorm_5', nn.BatchNorm2d(hidden_dim)),
('relu6_6', nn.ReLU6(inplace=True)),
# pw-linear
('conv2d_7', nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False)),
('bnorm_8', nn.BatchNorm2d(oup)),
('dq', torch.quantization.DeQuantStub())
]))
def forward(self, x):
if self.use_res_connect:
return x + self.conv(x)
else:
return self.conv(x)
下面是执行量化传递的函数:
def prepare_model(model):
model.train()
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model = torch.quantization.fuse_modules(
model,
[
# NOTE(aleksey): 'features' is the attr containing the non-head layers.
['features.in_conv.conv2d', 'features.in_conv.batchnorm2d'],
['features.inv_conv_1.conv.conv2d_1', 'features.inv_conv_1.conv.bnorm_2'],
['features.inv_conv_1.conv.conv2d_4', 'features.inv_conv_1.conv.bnorm_5'],
*[
*[[f'features.inv_conv_{i}.conv.conv2d_1',
f'features.inv_conv_{i}.conv.bnorm_2'] for i in range(2, 18)],
*[[f'features.inv_conv_{i}.conv.conv2d_4',
f'features.inv_conv_{i}.conv.bnorm_5'] for i in range(2, 18)],
*[[f'features.inv_conv_{i}.conv.conv2d_7',
f'features.inv_conv_{i}.conv.bnorm_8'] for i in range(2, 18)]
]
]
)
model = torch.quantization.prepare_qat(model)
return model
一旦我们定义并训练了我们的量化模型达到收敛，我们就可以将它传递到 TVM 优化引擎。这个过程的第一步是将计算图从追踪的 PyTorch 转换成 Relay:
import tvm
from tvm.contrib import graph_executor
import tvm.relay as relay
TARGET = "llvm -mcpu=skylake-avx512"
def get_tvm_model(traced_model, X_ex):
mod, params = relay.frontend.from_pytorch(
traced_model, input_infos=[('input0', X_ex.shape)]
)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=TARGET, params=params)
dev = tvm.device(TARGET, 0)
module = graph_executor.GraphModule(lib["default"](dev))
module.set_input("input0", X_ex)
module.run() # smoke test
# mod and params are IR structs used downstream.
# module is a Relay Python callable.
return mod, params, module
这个方法从调用 tvm.relay.frontend.from_pytorch开始。from_pytorch 需要两个东西: 跟踪的 PyTorch 模块(这里指traced_model)和解释模型输入形状的结构体。在这段代码中，X_ex 是从训练循环的 dataloader 中取样的一个示例批次，因此输入形状是从 X_ex.shape 得到的。
注意，输入有一个名称 input0。这个名称参数是必需的，因为 Relay 要为它的图输入命名。尽管 PyTorch 没有这样的概念，但 TVM 预期我们设置一个名称，不过它的实际值并不重要。
下一个调用 relay.build 实际上构造了 Relay 的计算图。它最重要的参数是 target; 这是运行此代码的硬件平台(以及目标)的字符串表示形式。尽可能明确地设置这个字符串来匹配你的目标平台非常重要，但不幸的是，我在文档中没有看到字符串参数列表。我使用 AWS 上的一个 c5.4xlarge 实例来运行这段代码，实例的芯片是 Intel Xeon Platinum 8000系列，因此 target 参数是:
llvm— 使用 llvm 编译器，因为这是一个 x86 芯片skylake— 这个芯片使用 Skylake 架构avx512— 这个芯片支持 AVX-512 扩展指令集
lib 和 mod 是指向 C(?) blobs 的指针，不能直接使用。Relay API 将 lib 包装在 GraphExecutor 中，创建了一个可以直接从 Python 调用的模块。
最后也是最重要的一步是调优:
def tune(mod, params, X_ex):
number = 10
repeat = 1
min_repeat_ms = 0
timeout = 10
# create a TVM runner
runner = autotvm.LocalRunner(
number=number,
repeat=repeat,
timeout=timeout,
min_repeat_ms=min_repeat_ms,
)
tuning_option = {
"tuner": "xgb",
"trials": 10,
"early_stopping": 100,
"measure_option": autotvm.measure_option(
builder=autotvm.LocalBuilder(build_func="default"), runner=runner
),
"tuning_records": "resnet-50-v2-autotuning.json",
}
tasks = autotvm.task.extract_from_program(
mod["main"], target=TARGET, params=params
)
for i, task in enumerate(tasks):
prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))
tuner_obj = XGBTuner(task, loss_type="rank")
tuner_obj.tune(
n_trial=min(tuning_option["trials"], len(task.config_space)),
early_stopping=tuning_option["early_stopping"],
measure_option=tuning_option["measure_option"],
callbacks=[
autotvm.callback.progress_bar(
tuning_option["trials"], prefix=prefix
),
autotvm.callback.log_to_file(tuning_option["tuning_records"]),
],
)
with autotvm.apply_history_best(tuning_option["tuning_records"]):
with tvm.transform.PassContext(opt_level=3, config={}):
lib = relay.build(mod, target=TARGET, params=params)
dev = tvm.device(str(TARGET), 0)
optimized_module = graph_executor.GraphModule(lib["default"](dev))
optimized_module.set_input("input0", X_ex)
optimized_module.run() # dry run test
return optimized_module
这段代码使用 XGBoost 库对 Relay 模型进行优化运行，在选定的时间约束条件下，为这个计算图找到尽可能接近最优的调度计划。这段代码中包含了大量的样板文件，但是你不必理解每一行代码所做的事情，就可以 get 到它的重点。
请注意，为了节省时间，我们执行的是一个只有10次测试运行的试验。对于生产用例，TVM 的应用 Python 入门指南推荐 CPU 运行1500次测试，GPU 运行3000次左右。
对结果模型进行基准测试
我记录了在 CPU 上这个模型的两个不同版本运行一批数据的时间，并计算了运行多次推理所需的平均时间。第一个是基准的 PyTorch 模型，没有量化和编译。第二个是完全优化的模型: 一个已经被量化，编译过的 MobileNet，并使用前面部分的代码进行调优。你可以在这里看到基准测试代码：https://github.com/spellml/examples/blob/master/tvm/scripts/test_mobilenet.py。以下是结果:
模型的编译版本的推理时间比基准模型快30倍以上！
事实上，值得注意的是，在 CPU 上编译的模型运行速度与 GPU 上的基准模型(g4dn.xlarge，NVIDIA T4实例)相当。因此，量化和模型编译带来的性能提升使得 CPU 和 GPU 的服务效率几乎一样，考虑到模型在优化之前的速度之慢，这一点非常显著。
请注意，并非所有这些性能提升都归功于 TVM，其中一部分来自于编译步骤之前应用于 PyTorch 模型的量化。
训练愉快！