前言
有人说当下的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模型的快速落地和普及,真正使一般大众都能玩起来,打破技术垄断。
备注
改变opset_version和量化也有可能提升性能,但在本例中均不明显,且量化后反而性能负提升。考虑到有人可能需要,在这里放两个链接:
改变opset_version
量化模型至int8