如果我只有一个奔腾CPU，怎么加速推理神经网络？

news2024/9/27 19:24:28

前言

有人说当下的AI热潮不过是算力堆砌的产物。现在层出不穷的各种大模型大训练集，使用复杂精致的技术在排行榜上不断刷新分数，这些人似乎忘了一件事情，AI模型最终是要落地的，是要用的，如果不能普及开去那和在象牙塔中精心雕刻雕塑又有什么区别呢？尽管现在有不少简单易用的框架都支持硬件加速，但gpu高昂的成本和厂家自我研发的诸如NPU之类的框架复杂度依然存在。所以本文将以一个天真可爱的想法展开，如果我就是个穷鬼，掏不起钱买显卡，只有一个五年前的Intel低端cpu，但我也想玩时下流行的style transfer GAN，该怎么办？

环境

首先迎来本文的主角1号，我们的style transfer模型xxx.onnx：
处于隐私保护，这里不能透露主角的姓名，只能通过输入输出来管中窥豹捏。
在这里插入图片描述
可以看到输入是1,?,?,3，输出是?,?,?,3，我们确认这是个nhwc的输入，和onnx的标准输入形状ncwh是不一样的捏，据说这点会影响性能推理，但实测下来并没有。此外，？号表示动态输入，说明这个模型的输入输出都没有固定形状。如果你还是听不懂的话，就是图片进模型图片出模型捏，像下面这样：

原图：
在这里插入图片描述
生成图：

然后是本文主角2号，咱们高贵的奔腾cpu：

最后是本文主角3号：
python3.10.6
Onnxruntime1.14.0
在这里插入图片描述

加速

1.基准测试

import onnxruntime as ort
options = ort.SessionOptions()
options.enable_profiling = True #记录性能开关
result = session.run(None, {x: img})[0] #运行推理
prof_file = session.end_profiling() #推理结束后结束日志

首先我们需要对推理速度进行一个基本测试，在对代码进行多段time.time()插入计时后，如果我们需要具体看ONNX每层计算的时间，则需要使用它提供的profiling功能，方法很简单，先把开关打开然后在推理结束后终止记录，就会在同一目录下生成onnxruntime_profile_XXXX-XX-XX_XX-XX-XX.json文件。结果如下：
在这里插入图片描述
可以看到瓶颈就是在推理这里，用了9.37s，那么到底是什么算子会如此耗时呢？我们用以下代码查看json文件：

import json

with open('onnxruntime_profile__2023-03-02_10-23-47.json', "r") as f:
    sess_time = json.load(f)
l = sorted(sess_time, key=lambda k: k['dur'], reverse=True)
per = 0
total = l[0]['dur']
for x in l[2:52]:
    print(x['dur'])#打印持续时间
    print(x['args'])#打印具体参数
    per+=x['dur']
print(f'Top50 total is {per/total*100:.2f}%')

这段代码能打印前五十[2:52]最耗时间的算子，并且给出占比。需要注意的是前两个指的是总共时间，所以这里把它排除了，结果如下：

在这里插入图片描述
注意到无一例外全都是可恶的Conv层，卷积是真的耗时啊。

然后我们再跑一次，这次看看硬件使用率如何：
在这里插入图片描述
在卷积层这么耗时的情况下, cpu竟然在偷懒！

2.解决CPU偷懒的问题
在前面那段代码上继续加一个选项，把使用线程数调到4，我这个cpu是2核4线程，那就把它吃满！狠狠的用看看有没有提升：

options.intra_op_num_threads = 4

现在已经吃满了：

在这里插入图片描述
效果如下：

在这里插入图片描述

提升效果明显啊，提升了1s左右。9.3->7.7。

3.关闭ONNXRUNTIME本身的优化

options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL

默认是全部开启优化的，我们把它关掉，发现性能反而提升了：

在这里插入图片描述
7.7->7.0，又快了0.7s。在经过后面GPU上的测试后发现该优化对gpu有一丁点的加速作用，对cpu来说反而是副作用。

4.固定尺寸+量化前预处理
这里假设我们已经对图片缩放到了1280x720分辨率。
先以1,720,1280,3的输入向量固定尺寸：

python -m onnxruntime.tools.make_dynamic_shape_fixed --input_name xxx --input_shape 1,720,1280,3

再进行量化前预处理，即尺寸推断和优化处理：

python -m onnxruntime.quantization.preprocess --input xxx.onnx --output xxx_preprocessed.onnx

再拿去跑一遍，发现又快了不少：
在这里插入图片描述
从7.0s快到了约6s，又提升1s。

总结

本例中，我们使用一些技巧成功使破烂cpu的推理速度加快了3s，得到了33%的性能提升，尽管提升看起来可能不够大，但在批量图片的生成过程中这就意味着节省几个小时。个人而言，我是非常希望所有的AI模型能以极低的代价和良好的性能为标准，推进AI模型的快速落地和普及，真正使一般大众都能玩起来，打破技术垄断。