1 转换的主要内容

在另外一篇博客中已经写了，只是在旭日3开发板上实操时用的是另外一个程序，故在此重新记录一下，详细转换过程参考如下链接

https://www.jianshu.com/p/eba19f6eba55

2 转换时主要程序

其中from nets.yolo import YoloBody可见参考链接

import onnxruntime
import torch
import numpy as np
import onnx

from nets.yolo import YoloBody


if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    input_shape = (416, 416)
    anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
    num_classes = 80        # 类别数要和下面model_path权重匹配
    net = YoloBody(anchors_mask, num_classes)
    model_path = './model_data/yolo_weights.pth'
    net.load_state_dict(torch.load(model_path, map_location=device))
    net = net.eval()
    # print(net)

    # input_shape = (416, 320)        # 导出onnx模型的输入尺寸，要和pytorch模型的输入尺寸一致
    dummy_input = torch.autograd.Variable(
        torch.randn(1, 3, input_shape[0], input_shape[1])
    )
    output_path = './model_data/yolov3_coco.onnx'       # onnx模型输出到哪里去
    output_names = None     # 本来应该有个响当当的名字，我不太会


    torch.onnx.export(
        net,
        dummy_input,
        output_path,
        verbose=True,
        keep_initializers_as_inputs=True,
        opset_version=11,       # 版本通常为10 or 11
        output_names=output_names,
    )

3 转换成onnx模型后进行验证

import onnxruntime
import torch
import numpy as np
import onnx

from nets.yolo import YoloBody


if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    input_shape = (416, 416)
    anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
    num_classes = 80        # 类别数要和下面model_path权重匹配
    net = YoloBody(anchors_mask, num_classes)
    model_path = './model_data/yolo_weights.pth'
    net.load_state_dict(torch.load(model_path, map_location=device))
    net = net.eval()
    # print(net)

    # input_shape = (416, 320)        # 导出onnx模型的输入尺寸，要和pytorch模型的输入尺寸一致
    dummy_input = torch.autograd.Variable(
        torch.randn(1, 3, input_shape[0], input_shape[1])
    )
    output_path = './model_data/yolov3_coco.onnx'       # onnx模型输出到哪里去
    output_names = None     # 本来应该有个响当当的名字，我不太会

    onnx_model = onnx.load('./model_data/yolov3_coco.onnx')
    onnx.checker.check_model(onnx_model)

    # # Input to the model
    x = torch.randn(size=(1, 3, input_shape[0], input_shape[1]), dtype=torch.float32).to(device)
    with torch.no_grad():
        torch_out = net(x)
    # print(type(torch_out))        # <class 'tuple'>，啥操作都干不了，元组里面是个张量
    # print(type(torch_out[0]))     # <class 'torch.Tensor'>
    torch_out = torch_out[0]        # 因此这儿需要把内部张量取出来
    # print(torch_out.shape)          # torch.Size([1, 255, 13, 13])

    ort_session = onnxruntime.InferenceSession("./model_data/yolov3_coco.onnx")  # 初始化模型
    ort_inputs = {ort_session.get_inputs()[0].name: x.numpy()}  # 初始化数据，注意这儿的x是上面的输入数据x，后期应该是img
    ort_outs = ort_session.run(None, ort_inputs)  # 推理得到输出
    # print(type(ort_outs))       # <class 'list'>，里面是个numpy矩阵
    # print(type(ort_outs[0]))    # <class 'numpy.ndarray'>
    ort_outs = ort_outs[0]  # 因此这儿需要把内部numpy矩阵取出来

    # print(torch_out.numpy().shape)      # (1, 255, 13, 13)
    # print(ort_outs.shape)  # (1, 255, 13, 13)

    # 比较实际值与期望值的差异，通过没啥事，不通过引发AssertionError
    np.testing.assert_allclose(torch_out.numpy(), ort_outs, rtol=1e-03, atol=1e-05)

    print('torch_out:', torch_out.numpy())
    print('ort_outs:', ort_outs)

4 文件所在位置

文件所在位置