通过解读yolov5_gpu_optimization学习如何使用onnx_surgon

onnx实战一: 解析yolov5 gpu的onnx优化案例:
这是一个英伟达的仓库, 这个仓库的做法就是通过用gs对onnx进行修改减少算子然后最后使用TensorRT插件实现算子， 左边是优化过的, 右边是原版的。 通过这个案例理解原版的onnx的导出流程然后我们看英伟达是怎么拿gs来优化这个onnx

原版的export_onnx函数

先看torch.onnx.export函数的参数解释：

1. model: 要导出的PyTorch模型, 在工程中这里输入的是训练好的pt文件

2. im: 这里对应torch.onnx.export的args, 这个是用作模型输入的示例张量。这帮助ONNX确定输入的形状和类型。

3. f: 输出ONNX模型的文件名或文件对象, 用来指定导出模型的路径和文件名。

4. verbose (默认为 False): 如果设置为 True，则会打印出模型导出时的详细日志。

5. opset_version: 导出的ONNX模型的操作集版本。不同的版本可能支持不同的操作。

6. training:

   • torch.onnx.TrainingMode.TRAINING: 表示模型处于训练模式。
   • torch.onnx.TrainingMode.EVAL: 表示模型处于评估模式。

7. do_constant_folding (默认为 True): 当设置为 True，导出过程中会尝试简化模型，将常量子图折叠为一个常量节点。

8. input_names: 为模型的输入提供名称, 参数规定是数组

9. output_names: 为模型的输出提供名称, 参数规定是数组

10. dynamic_axes: 为模型的输入/输出定义动态轴。对于那些维度在推理时可能会发生变化的情况（例如，批处理大小），此参数允许指定哪些轴是动态的。这里images是输入, 本来是1x3x640x640, 这里通过指定把0, 2, 3维度变成了动态轴的输入, 第二个维度是3这个还是固定的。如果使用动态, 可以输入任意数量和任意大小的图片而不是规定的单张640x640

• 'images': 对应的张量名称。
  • 0: 'batch': 表示第0个维度（即批处理维度）是动态的，并命名为’batch’。
  • 2: 'height': 表示第2个维度（即图像的高度）是动态的。
  • 3: 'width': 表示第3个维度（即图像的宽度）是动态的。
• 'output': 对应的张量名称。
  • 0: 'batch': 表示第0个维度（即批处理维度）是动态的。
  • 1: 'anchors': 表示第1个维度是动态的。

11. dynamic (没有在给定的函数调用中明确给出，但可以从上下文推断):

• True: 如果你想让某些轴动态，你可以设置此参数为True。
• False: 表示不使用动态轴。

导出了onnx之后开始做onnxsim

1. model_onnx, check = onnxsim.simplify(...):使用onnxsim的simplify方法简化模型。它返回简化后的onnx模型和一个布尔值check，表示简化是否成功。

2. 在对动态输入的onnx导出的时候, dynamic_input_shape=dynamic是不够的，还要把输入给他，让onnxsim更加谨慎的优化onnx, 确保满足我们给他的输出，所以这里多了一个input_shapes={'images': list(im.shape)} if dynamic else None

def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
    # YOLOv5 ONNX export
    try:
        check_requirements(('onnx',))
        import onnx

        LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
        f = file.with_suffix('.onnx')

        torch.onnx.export(
            model,
            im,
            f,
            verbose=False,
            opset_version=opset,
            training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
            do_constant_folding=not train,
            input_names=['images'],
            output_names=['output'],
            dynamic_axes={
   
   
                'images': {
   
   
                    0: 'batch',
                    2: 'height',
                    3: 'width'},  # shape(1,3,640,640)
                'output': {
   
   
                    0: 'batch',
                    1: 'anchors'}  # shape(1,25200,85)
            } if dynamic else None)

        # Checks
        model_onnx = onnx.load(f)  # load onnx model
        onnx.checker.check_model(model_onnx)  # check onnx model

        # Metadata
        d = {
   
   'stride': int(max(model.stride)), 'names': model.names}
        for k, v in d.items<