将yolov3.pth转换为onnx并进行推理

前言

      最近训练了一个yolov3检测模型，需要将保存好的.pth文件转换为onnx并且写一个inference过程。于是学习了一下onnx的知识并记录下来。

用到的工具

onnx

安装方法：pip install onnx

onnxruntime

安装方法：pip install onnxruntime

pth-->onnx

1.加载训练好的模型和模型参数

2.使用model.eval()

3.创建一个虚拟输入，输入的形状和你训练的模型接受的输入形状要求一致

4.使用torch.onnx.export进行转换，具体操作为：torch.onnx.export(model,input,"保存路径")(注：其他参数可以自行了解)

5.转换完成

推理

导入onnx和onnxruntime

model=onnx.load("onnx文件路径")

加载完成后可以使用onnx.checker.check_model(model)来检查有没有问题。一般来说是不会出问题的。

然后加载输入，并且对输入做前处理，比如我这里是对图片进行检测，所有我做的前处理包括读入图片，将图片处理到需要的shape。

创建onnxruntime session：sess=onnxruntime.InferenceSession("onnx文件路径")

input_name=sess.get_inputs()

result=sess.run(None,{input_name[0].name:img.numpy()}),此处的img即是之前进行过前处理的图像，result就是这张图像经过检测模型得到的结果，之后再经过nms和后处理就可以得到想要的结果了！

后记

目前只是对onnx有了一个初步的理解，可能很多地方还有不足的地方，希望大家可以多多包涵，有问题的话也欢迎交流，之后会继续补充！