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 文件所在位置
文件所在位置
