当前位置: 首页 > 工具软件 > node-tiny > 使用案例 >

【yolov4-tiny】二、pytorch->onnx->caffe

管景天
2023-12-01

一、Pytorch->onnx

使用https://github.com/bubbliiiing/yolov4-tiny-pytorch,训练,导出onnx

import torch.onnx
import onnx
from onnxsim import simplify

yolo = YOLO()
model = yolo.net
x = torch.randn(1,3,416,416)
out_name = "yolov4-tiny_pytorch.onnx"
torch.onnx.export(model, x, out_name, export_params=True, opset_version=9)

model = onnx.load(out_name)
model_simp, check = simplify(model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, out_name)

二、onnx->caffe

1、搭环境

由于onnxruntime不支持Python2,caffe也得是Python3环境下的

从docker里找了个:

docker pull mapler/caffe-py3:cpu

顺便记录一下docker里怎么科学上网:

pip3 install pysocks
export all_proxy="socks5://(ip):(port)"

安装onnx:

pip3 install onnx==1.6.0
pip3 install onnxruntime==1.1.0

扩展caffe算子:主要是Upsample,参考笔者之前的文章【yolov4-tiny】darknet->caffe

2、转模型

使用https://github.com/xxradon/ONNXToCaffe的代码

python3 convertCaffe.py ./model/yolov4-tiny_pytorch.onnx ./model/yolov4-tiny_pytorch.prototxt ./model/yolov4-tiny_pytorch.caffemodel

注释掉了convertCaffe.py里的一行不支持算子的代码,这里是权重转换,没有权重的算子当然不支持,能继续跑就行了

for id, node in enumerate(graph.nodes):
    node_name = node.name
    op_type = node.op_type
    inputs = node.inputs
    inputs_tensor = node.input_tensors
    input_non_exist_flag = False
    if op_type not in wlr._ONNX_NODE_REGISTRY:
        # err.unsupported_op(node)
        continue
    converter_fn = wlr._ONNX_NODE_REGISTRY[op_type]
    converter_fn(net, node, graph, err)
 类似资料: