使用 pnnx 优化 swin transformer 6倍性能

0x0 先贴 pnnx 优化效果

• torch 2.0
• nvidia T4
• fp16 推理

0x1 使用 torch.compile 优化 swin_t

import torch
import torchvision

torch.set_grad_enabled(False)

torch.manual_seed(0)
x = torch.rand(1, 3, 224, 224).half().cuda()

# torchvision model inference
net = torchvision.models.swin_t(pretrained=True)
net.eval().half().cuda()

a = net(x)

net_compiled = torch.compile(net, mode='max-autotune')

a2 = net_compiled(x)

0x2 使用 pnnx 优化 swin_t

下载二进制程序 pnnx

https://github.com/pnnx/pnnx/releasesgithub.com/pnnx/pnnx/releases

# export torchscript
mod = torch.jit.trace(net, x)
mod.save("swin_t.pt")

# torchscript to pnnx
import os
os.system("./pnnx swin_t.pt inputshape=[1,3,224,224]")

# pnnx inference
import swin_t_pnnx
pnnx_net = swin_t_pnnx.Model()
pnnx_net.eval().half().cuda()

b = pnnx_net(x)

pnnx_net_compiled = torch.compile(pnnx_net, mode='max-autotune')

b2 = pnnx_net_compiled(x)

对比结果可得知，对应相同输入 x，输出 a2 和 b2 仅有小数点末尾精度差异，因此模型转换正常可用

0x3 跑1000次测速

import time

# warmup
for _ in range(100):
    a = net(x)
torch.cuda.synchronize()
start = time.time()
for _ in range(1000):
    a = net(x)
torch.cuda.synchronize()
print('time = ', time.time() - start)

# same for net_compiled pnnx_net and pnnx_net_compiled
# ... omitted

0x4 pnnx 优化原理

• 把 swin_t 的 attention 实现自动替换为 nn.MultiheadAttention
• 自动常量折叠，如 shape，relative_position_bias 等等
• 其他常规优化，如表达式合并，死代码消除，conv-bn合并等等
• gelu 遵从原始模型设置，没有使用 tanh 近似，还有点优化空间

0x5 一些 FAQ

• pnnx 是什么？

PyTorch Neural Network Exchange。实质是一种 pytorch 模型的表达方法，看这里的介绍

nihui：PNNX: PyTorch Neural Network Exchange

https://zhuanlan.zhihu.com/p/427620428

• 如何安装 pnnx？

推荐下载编译好的 pnnx 二进制包，这个 pnnx 程序不依赖任何第三方库和 python 环境，portable，支持 Windows/Linux/MacOS

https://github.com/pnnx/pnnx/releasesgithub.com/pnnx/pnnx/releases

也可以自己编译，源代码在

https://github.com/Tencent/ncnn/tree/master/tools/pnnx

• 为什么使用 pnnx 而不是 onnx / onnx-runtime？

pnnx 解决了许多 onnx 暂时无法解决的问题

1. >2G的大模型导出和优化
2. 支持更多的 pytorch 算子
3. 支持自定义 op
4. 针对 pytorch 模型有更多优化
5. 不添加胶水op和冗余操作
6. .....等

经过优化后的模型直接兼容 pytorch，方便后续再训练和部署等

• pnnx 怎么用？

如本例，先 trace 保存 torchscript，连带输入 shape 传入 pnnx 命令行参数，自动转换出对应的 pnnx 模型文件和优化后模型的 py 代码

• pnnx 支持哪些模型？

pnnx 目前已支持绝大部分常见算子和模型，sherpa-ncnn，vits-ncnn，stable-diffusion-ncnn 等开源项目都使用 pnnx 做模型优化和转换。如果遇到不支持的 op，请向 ncnn github 发 issue 或加 QQ 群反馈

• pnnx 加速效果如何？

取决于 torch 底层的实现效率。最新 pytorch-2.0 中大幅优化 Attention 相关算子性能，因此在 swin transformer 中效果明显

• pnnx 优化的模型有额外依赖吗？

pnnx 优化后，只依赖import torch，不再需要 import 其他库，精简依赖关系，方便部署和后续操作

• pnnx 优化的模型支持转 onnx 吗？

pnnx 优化后是标准 torch Module，可直接导出 onnx，但 onnx 会再次拆散 MultiheadAttention，使得模型变复杂和降速（onnx坏！

• pnnx 支持其他框架吗，比如 paddle onnx 等？

pnnx 目前只支持 pytorch 的模型，其他框架推荐使用 onnx2torch 等工具先转为 pytorch 模型再使用

• pnnx 有交流群吗，遇到问题怎么办？

欢迎加 ncnn 技术交流群，群号和入群暗号见 ncnn github readme