YOLOV3 转ONNX，于移动端应用

@[TOC]YOLOV3 转ONNX，于移动端应用

GitHub代码

最近整理了在不使用GPU情况下，于x86下进行神经网络前向推理加速的代码，具体链接可见上方。
这里的代码是目标检测相关的，基于pytorch 的 yolov3，因为pytorch 无法直接转换至openvino对应模型，故模型转换时，采用了如下步骤，pytorch >> onnx >> openvino。
onnx 结构可以视作与Tensorflow、Caffe类似的静态图，故在转换时将yolov3中用到的锚框直接作为常量嵌入至onnx模型中，具体转换代码如下，转换前记住pip install onnx欧：

#-*- coding:utf-8 -*-

import os
import sys
import onnx
import onnxruntime
import numpy as np
import torch
from models import Darknet

def transform_to_onnx(cfg_file, weight_file, batch_size, in_h, in_w):
    model = Darknet(cfg_file)
    pre_dict = torch.load(weight_file, map_location=torch.device('cpu'))
    model.load_state_dict(pre_dict['model'])
    x = torch.ones((batch_size, 3, in_h, in_w), requires_grad=True)*120 /255.0
    onnx_file_name = 'model/yolov3.onnx'
    
    torch.onnx.export(model,
                      x,
                      onnx_file_name,
                      #export_params=True,
                      #opset_version=11,
                      #do_constant_folding=True,
                      input_names=['input'], output_names=['output1'])
                      #dynamic_axes=None)
    print('Onnx model exporting done')
    return onnx_file_name, x

def main(cfg_file, weight_file, batch_size, in_h, in_w):
    onnx_path_demo, x = transform_to_onnx(cfg_file, weight_file, 1, in_h, in_w)
    #session = onnxruntime.InferenceSession(onnx_path_demo)
    #output = session.run(['output1'], {'input':x.detach().numpy()})


if __name__ == "__main__":
    cfg_file = 'cfg/yolov3-tiny.cfg'
    weight_file = 'model/last.pt'
    batch_size = 1
    in_h = 416
    in_w = 416
    main(cfg_file, weight_file, batch_size, in_h, in_w)

openvino 前向推理等相关代码，详见github。如果帮到了你，请给我们点赞呦。