三. TensorRT基础入门-快速分析开源代码并导出onnx

前言

自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习课程第三章—TensorRT 基础入门，一起来学习 Swin Transformer 的 ONNX 导出

课程大纲可以看下面的思维导图

0. 简述

本小节目标：以 Swin Transformer 为例学习快速导出 onnx 并分析 onnx 的方法

这节我们学习第三章节第八小节—快速分析开源代码并导出 ONNX，我们以 Swin Transformer 为例来讲解怎么快速导出一个 ONNX，并分析它里面的架构是什么样的

Swin Transformer 官方并没有提供一个 ONNX 的导出方法，但是我们可以按照之前学习的方式把它的 ONNX 导出来。同时它也是一个比较好的案例，这是因为当我们导出它的 ONNX 时会出错，主要涉及一些算子的兼容性问题，因此我们还可以从中学习当这些算子出现错误的时候我们如何解决它从而正确导出 ONNX

我们通过本次课程来讲解遇到不兼容的算子时该如何导出 ONNX，大家以后如果有类似的问题可以参考下

1. 执行一下我们的python程序

源代码获取地址：https://github.com/kalfazed/tensorrt_starter

Swin Transformer 项目地址：https://github.com/microsoft/Swin-Transformer

这个小节的案例主要是 3.6-export-onnx-from-oss，如下所示：

其中：

• models：该文件夹用于存放导出的 ONNX
• Swin-Transformer：该文件夹是 Swin-Transformer 的源码
• weights：该文件夹用于存放 Swin-Transformer 模型的权重
• src：该文件夹存放导出 ONNX 的 export.py 脚本文件

下面是导出 ONNX 的部分结果：

2. 转换swin-tiny时候出现的不兼容op的解决方案

我们先配置 Swin Transformer 的环境，可以在 Swin-Transformer/get_started.md 中找到如下的安装教程：

值得注意的是博主并没有配置这些环境，主要是因为这个环境是训练测试 Swin Transformer 时所使用的环境，其中的很多依赖库其实我们并不需要

我们只需要导出它的 ONNX 就行，因此博主还是拿的自己当前的虚拟环境进行后续 ONNX 的导出，导出过程中如果遇到需要安装的库再一个个安装就行

下面我们就来看看如何导出 ONNX，我们先进 main.py 看看这个文件中所做的事情，主要是观察有没有一些奇怪的操作可能影响我们的 ONNX 导出，另外看一看模型的构建是如何做的，它的部分代码如下所示：

def main(config):
    dataset_train, dataset_val, data_loader_train, data_loader_val, mixup_fn = build_loader(config)

    logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}")
    model = build_model(config)
    logger.info(str(model))

    ...

可以看到我们模型的构建主要是通过 build_model 这个函数来完成的，拿到 model 之后其实我们就可以来进行 ONNX 的导出了

这里韩君老师写好了一个 ONNX 导出的脚本文件，完整代码如下所示：

# --------------------------------------------------------
# Swin Transformer
# Copyright (c) 2021 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ze Liu
# --------------------------------------------------------

import argparse

import torch

from config import get_config
from models import build_model


import onnx
import onnxsim

def parse_option():
    parser = argparse.ArgumentParser('Swin Transformer training and evaluation script', add_help=False)
    parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file', )
    parser.add_argument(
        "--opts",
        help="Modify config options by adding 'KEY VALUE' pairs. ",
        default=None,
        nargs='+',
    )

    # easy config modification
    parser.add_argument('--batch-size', type=int, help="batch size for single GPU")
    parser.add_argument('--data-path', type=str, help='path to dataset')
    parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
    parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'],
                        help='no: no cache, '
                             'full: cache all data, '
                             'part: sharding the dataset into nonoverlapping pieces and only cache one piece')
    parser.add_argument('--pretrained',
                        help='pretrained weight from checkpoint, could be imagenet22k pretrained weight')
    parser.add_argument('--resume', help='resume from checkpoint')
    parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps")
    parser.add_argument('--use-checkpoint', action='store_true',
                        help="whether to use gradient checkpointing to save memory")
    parser.add_argument('--disable_amp', action='store_true', help='Disable pytorch amp')
    parser.add_argument('--amp-opt-level', type=str, choices=['O0', 'O1', 'O2'],
                        help='mixed precision opt level, if O0, no amp is used (deprecated!)')
    parser.add_argument('--output', default='output', type=str, metavar='PATH',
                        help='root of output folder, the full path is <output>/<model_name>/<tag> (default: output)')
    parser.add_argument('--tag', help='tag of experiment')
    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
    parser.add_argument('--throughput', action='store_true', help='Test throughput only')

    # distributed training
    parser.add_argument("--local_rank", type=int, required=True, help='local rank for DistributedDataParallel')

    # for acceleration
    parser.add_argument('--fused_window_process', action='store_true',
                        help='Fused window shift & window partition, similar for reversed part.')
    parser.add_argument('--fused_layernorm', action='store_true', help='Use fused layernorm.')
    ## overwrite optimizer in config (*.yaml) if specified, e.g., fused_adam/fused_lamb
    parser.add_argument('--optim', type=str,
                        help='overwrite optimizer if provided, can be adamw/sgd/fused_adam/fused_lamb.')

    args, unparsed = parser.parse_known_args()

    # import os
    # if 'LOCAL_RANK' not in os.environ:
    #     os.environ['LOCAL_RANK'] = '0'

    config = get_config(args)

    return args, config

class Model(torch.nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        self.softmax  = torch.nn.Softmax()

    def forward(self, x):
        x = self.backbone(x)
        x = self.softmax(x)
        return x

def export_norm_onnx(model, file, input):
    torch.onnx.export(
        model         = model, 
        args          = (input,),
        f             = file,
        input_names   = ["input0"],
        output_names  = ["output0"],
        opset_version = 9)

    print("Finished normal onnx export")

    model_onnx = onnx.load(file)

    # 检查导入的onnx model
    onnx.checker.check_model(model_onnx)

    # 使用onnx-simplifier来进行onnx的简化。
    print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")
    model_onnx, check = onnxsim.simplify(model_onnx)
    assert check, "assert check failed"
    onnx.save(model_onnx, file)

def main(config):
    backbone = build_model(config)
    model = Model(backbone=backbone)

    input = torch.rand(1, 3, 224, 224)

    model.eval()
    export_norm_onnx(model, "../models/swin-tiny-after-simplify-opset9.onnx", input)
    # export_norm_onnx(model, "../models/swin-tiny-after-simplify-opset12.onnx", input)
    # export_norm_onnx(model, "../models/swin-tiny-after-simplify-opset17.onnx", input)

# python export.py --eval --cfg configs/swin/swin_tiny_patch4_window7_224.yaml --resume weights/swin_tiny_patch4_window7_224.pth --data-path data --local_rank 0
if __name__ == '__main__':
    args, config = parse_option()
    main(config)

这个脚本实现了从命令行解析配置参数，构建 Swin Transformer 模型，并导出为 ONNX 格式，具体步骤如下：(from ChatGPT)

• 1. 解析命令行参数：使用 argparse 解析各种命令行参数，如配置文件路径、批量大小、数据路径等，并生成配置对象。
• 2. 定义模型类：定义一个包含主干网络和 Softmax 层的 Model 类，继承自 torch.nn.Module，用于封装 Swin Transformer 模型。
• 3. 导出模型为 ONNX 格式：定义 export_norm_onnx 函数，使用 torch.onnx.export 导出模型到指定文件，检查导出的 ONNX 模型并使用 onnxsim.simplify 进行简化，最后保存简化后的模型。
• 4. 主函数：在 main 函数中，构建模型的主干部分并封装到自定义的 Model 类中，生成一个随机输入张量，将模型设置为评估模式，并调用 export_norm_onnx 导出模型。
• 5. 程序入口：在脚本入口处，调用 parse_option 解析命令行参数，生成配置对象后，调用 main 函数执行上述主要流程。

这个脚本的目的是将训练好的 Swin Transformer 模型转换为 ONNX 格式，并简化模型以便于在其他深度学习框架中进行高效推理。

我们先看 opset=9 时的导出，导出指令如下：

python export.py --eval --cfg configs/swin/swin_tiny_patch4_window7_224.yaml --resume weights/swin_tiny_patch4_window7_224.pth --data-path data --local_rank 0

执行后你可能会遇到如下问题：

这个错误表明在导出 ONNX 模型时，slice 操作的步长不等于 1，而当前的 ONNX 导出工具不支持这种步长设置。Swin Transformer 可能在某些操作中使用了步长不等于 1 的 slice 操作。可以修改 swin transformer 的 slice 操作也可以尝试更新 PyTorch 到最新版本，有可能该问题已经在新版本中得到解决

博主选择换一个高版本的 Pytorch 虚拟环境，目前的 torch==1.12.1，新的 torch ==2.0.1，再次执行上述导出代码，你可能会遇到如下问题：

这个错误表明在 update_config 函数中尝试访问环境变量 LOCAL_RANK 时，该变量未被设置。我们可以手动设置，新增代码如下所示：

# export.py 64行

def parse_option():
    ...

    args, unparsed = parser.parse_known_args()
    
    import os
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = '0'
        
    config = get_config(args)

再次执行就后博主没有出现老师存在的 operator roll 算子兼容的问题，而是依旧出现了 slice 的问题，如下图所示：

大家可能会在 swin transformer 导出的时候会出现如下问题：

反正 opset9 导出时总是会出现这样或那样的问题，我们再来看看修改为 opset12 导出时会怎样，如下图所示：

导出的 ONNX 如下图所示：

可以看到博主成功导出了，和韩君老师的结果又不一样，老师导出时会依旧出现 roll 算子问题，如下图所示：

如果出现和老师一样的 roll 算子问题，我们来看如何解决，首先去官网看下：

可以看到 roll 算子并没有一个实现，下面我们就要进入代码中看下 roll 在哪里进行了使用：

这个就很麻烦，可以看到 pytorch 中有 roll 这个算子但是 onnx 中没有，这样的话我们就需要自己在 onnx 中进行 roll 算子的实现了：

我们进入 https://github.com/pytorch/pytorch/blob/main/torch/onnx/symbolic_opset12.py 中看下，这个文件是 torch 导出 onnx 时对 opset12 算子的支持：

可以看到这里面也没有 roll 的支持，但是有一些其余的算子如 einsum、dropout 等的支持，我们可以在自己的虚拟环境中找到 symbolic_opset12.py 文件，然后添加 roll 算子的支持就行

博主的文件完整路径是：E:\Anaconda\envs\zmq\Lib\site-packages\torch\onnx\symbolic_opset12.py

在 symbolic_opset12.py 添加 pytorch2onnx 的 roll 算子实现，如下图所示：

大家修改后可以再次尝试下 ONNX 能否导出，博主这边之前就已经能正常导出，所以就没在这里演示了

3. swin-tiny的onnx分析

值得注意的是 opset9 即便把 Roll 的问题解决还会出现其他算子不兼容问题，所以直接在 opset12 上进行解决，之前我们导出的 ONNX 是利用 onnx-simplifier 简化过后的，如果不简化的话导出的 ONNX 如下图所示：

可以看到不简化的话整体结构非常复杂，看起来很难看，这个时候我们就可以利用 onnx-simplifier 进行简化

onnxsim 对 LN 的简化如下图所示：

onnxsim 对 attention 的简化如下图所示：

4. 根据onnx中的Proto信息，修改onnx

这里临时补充下 LayerNormalization 的理论推导，LN 和 BN 的计算其实是差不多的，两者都是在做 normalization 只不过是在不同的维度上，具体实现如下：

step1：mean 和 std 的计算
$${\mu }^l=\frac{1}{H}\sum _{i=1}^H{a}_i^l{\sigma }^l=\sqrt{\frac{1}{H}\sum _{i=1}^H(a_i^l-{\mu }^l{)}^2}$$
step2：normalization
$${\hat{\mathbf{a}}}^l=\frac{{\mathbf{a}}^l-{\mu }^l}{\sqrt{({\sigma }^l{)}^2+ϵ}}$$
step3：创建一个可以学习的 mean 和 var，最后再接一个激活函数
$${\mathbf{h}}^l=f({\mathbf{g}}^l\odot {\hat{\mathbf{a}}}^l+{\mathbf{b}}^l)$$
总结下来：
$$\mathbf{h}=f(\frac{\mathbf{g}}{\sqrt{{\sigma }^2+ϵ}}\odot (\mathbf{a}-\mu )+\mathbf{b})$$
Transformer 中 LN 的处理如下图所示：

LN 所需要的计算时间如下图所示：

计算时间大概是 165us，主要卡在了 reshape 和 pointwise 上面了

另外我们之前说过在 opset=17 以后就开始支持 LayerNormalization 算子，另外 TensorRT-8.6 版本以后也直接开始支持 LayerNormalization 了，如下图所示：

在代码中如果我们将 opset 修改为 17 按理来说导出的 onnx 应该更为简洁，LayerNormalization 应该是一个完整的算子了，如下图所示：

5. TensorRT导出

导出了 onnx 不是重点，我们需要将这些 onnx 导出为 TensorRT 模型并查看性能，我们有以下几种方式导出：

• trtexec 命令行（快速测试推理引擎）
• TensorRT python API（大量的单元测试）
• TensorRT C++ API（结合其他的前处理后处理进行底层部署）

总结

本次课程我们主要学习了 Swin Transformer 的 ONNX 导出，在导出的过程中可能会出现一些节点算子的兼容性问题，我们可以自己写符号函数来支持不兼容的算子，另外 LayerNormalization 算子在高版本的 opset 已经支持整个算子完整导出。值得注意的是 ONNX 导出只是一个开始，我们需要利用导出的 ONNX 生成 TensorRT 的 Engine 并查看分析性能。

OK，以上就是第 8 小节有关快速分析开源代码并导出 onnx 的全部内容了，下节我们来学习 trtexec 工具的使用，敬请期待😄

参考

• https://pytorch.org/docs/stable/generated/torch.roll.html
• https://github.com/onnx/onnx/blob/main/docs/Operators.md
• https://github.com/pytorch/pytorch/blob/main/torch/onnx/symbolic_opset12.py
• https://github.com/onnx/onnx-tensorrt/blob/main/docs/Changelog.md