yolo模型转换:pytorch -> onnx -> caffe

第一步:pytorch转onnx(pytorch版,yolov3-9.0开始提供脚本export.py)

(1)设置onnx算子版本(按需)

  修改代码:
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
  为:
torch.onnx.export(model, img, f, verbose=True, opset_version=9, input_names=['images'], do_constant_folding=True,

(2)导出onnx

python3 ./models/export.py --weights ./weights/yolov3.pt

(3)onnx模型简化,项目地址: https://github.com/daquexian/onnx-simplifier,否则下一步报错TypeError: ONNX node of type Shape is not supported.

pip install onnx-simplifier

python3 -m onnxsim ./weights/yolov3.onnx ./weights/yolov3-si.onnx

 第二步:onnx转caffe

git clone https://github.com/xxradon/ONNXToCaffe.git

cd ONNXToCaffe

python3 convertCaffe.py ./model/yolov3-si.onnx ./model/yolov3-si.prototxt ./model/yolov3-si.caffemodel

报错一:

Traceback (most recent call last):
  File "convertCaffe.py", line 122, in 
    convertToCaffe(graph, opset_version, prototxt_path, caffemodel_path)
  File "convertCaffe.py", line 67, in convertToCaffe
    layer = converter_fn(node,graph,err)
  File "/home/ubuntu/ONNXToCaffe-master/onnx2caffe/_operators.py", line 257, in _convert_Reshape
    return err.unsupported_op_configuration(node, "Reshape dimention number shall be 2 or 4")
  File "/home/ubuntu/ONNXToCaffe-master/onnx2caffe/_error_utils.py", line 44, in unsupported_op_configuration
    raise TypeError(
TypeError: Error while converting op of type: Reshape. Error message: Reshape dimention number shall be 2 or 4

修改_operators.py(yolo的数据shape是5维数组)

vim ./onnx2caffe/_operators.py

  修改代码:
elif len(shape) == 4:
  为
elif len(shape) == 4 or len(shape) == 5:

报错二:AttributeError: permute_param 或者 AttributeError: Upsample_param

Traceback (most recent call last):
  File "convertCaffe.py", line 122, in 
    convertToCaffe(graph, opset_version, prototxt_path, caffemodel_path)
  File "convertCaffe.py", line 79, in convertToCaffe
    layers[id] = layer._to_proto()
  File "/home/ubuntu/ONNXToCaffe-master/MyCaffe.py", line 100, in _to_proto
    assign_proto(layer, k, v)
  File "/home/ubuntu/ONNXToCaffe-master/MyCaffe.py", line 29, in assign_proto
    is_repeated_field = hasattr(getattr(proto, name), 'extend')
AttributeError: permute_param

添加Upsample层、permute层,重新编译caffe,项目地址 :https://github.com/jnulzl/caffe_plus

git clone https://github.com/jnulzl/caffe_plus.git

cp caffe_plus/include/caffe/layers/upsample_layer.hpp caffe/include/caffe/layers/
cp caffe_plus/src/caffe/layers/upsample_layer.cpp caffe/src/caffe/layers/
cp caffe_plus/src/caffe/layers/upsample_layer.cu caffe/src/caffe/layers/
cp caffe_plus/include/caffe/layers/permute_layer.hpp caffe/include/caffe/layers/
cp caffe_plus/src/caffe/layers/permute_layer.cpp caffe/src/caffe/layers/
cp caffe_plus/src/caffe/layers/permute_layer.cu caffe/src/caffe/layers/

# 修改caffe.proto文件

vim caffe/src/caffe/proto/caffe.proto

  在optional WindowDataParameter window_data_param = 129;(约第423行)后增加代码:

optional PermuteParameter permute_param = 150;
optional UpsampleParameter upsample_param = 151;

  在末尾增加代码:
message PermuteParameter {
  // The new orders of the axes of data. Notice it should be with
  // in the same range as the input data, and it starts from 0.
  // Do not provide repeated order.
  repeated uint32 order = 1;
}

message UpsampleParameter {		
	optional int32 height = 1 [default = 32];
	optional int32 width = 2 [default = 32];
	optional int32 height_scale = 3 [default = 2];
	optional int32 width_scale = 4 [default = 2];
	enum UpsampleOp {
		NEAREST = 0;
		BILINEAR = 1;
	}
	optional UpsampleOp mode = 5 [default = BILINEAR];
}

# 重新编译

cd caffe
make all -j8
make pycaffe -j8

相关文章:一文带你熟悉Pytorch->Caffe->om模型转换流程 ,海思开发:yolo v5s :pytorch->onnx->caffe->nnie

相关项目:https://github.com/Wulingtian/yolov5_onnx2caffe

你可能感兴趣的:(烂笔头,pytorch,caffe)