torchvision onnx 模型导出

本文是将torchvision包中的模型导出为ONNX，然后使用 Netron 可视化，从而可以结合Pytorch源码学习网络结构。

import torch
import torchvision

def export_onnx(model, im, file, opset, train, dynamic, simplify):
    # ONNX export
    try:
        import onnx

        torch.onnx.export(model, im, file, 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(file)  # load onnx model
        onnx.checker.check_model(model_onnx)  # check onnx model
         # Simplify
        if simplify:
            try:
                import onnxsim

                model_onnx, check = onnxsim.simplify(
                    model_onnx,
                    dynamic_input_shape=dynamic,
                    input_shapes={'images': list(im.shape)} if dynamic else None)
                assert check, 'assert check failed'
                onnx.save(model_onnx, file)
            except Exception as e:
                print(f'{file} simplifier failure: {e}')
        print(f'{file} export success, saved as {file}')
        return file
    except Exception as e:
        print("导出失败:", e)

if __name__ == "__main__":
    # model = torchvision.models.squeezenet1_0(pretrained=True).cuda()
    model = torchvision.models.squeezenet1_1(pretrained=True).cuda()
    batch_size = 16

    im = torch.zeros(batch_size, 3, 224, 224).to('cuda')  # image size(batch_size,3,224,224) BCHW iDetection
    file = "squeezenet1_1.onnx"

    try:
        import tensorrt as trt
        if trt.__version__[0] == '7':  
            print(trt.__version__[0])
            export_onnx(model, im, file, 12, train = False, dynamic = False, simplify=True)  # opset 12
        else:  # TensorRT >= 8
            print(trt.__version__[0])
            export_onnx(model, im, file, 13, train = False, dynamic = False, simplify=True)  # opset 13

    except Exception as e:
        print("导出失败")

本文导出的是 squeezenet 模型。

后续会介绍该模型！