使用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)
由于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
使用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)