|
pytorch转caffe模型的实战操作
总体思路是:pytorch ---> onnx ---> caffe
pytorch转onnx很容易:torch.onnx.export
http://www.ai111.vip/thread-1077-1-1.html
onnx转caffe :https://github.com/MTlab/onnx2caffe
遇到的问题及解决:
1、onnx转caffe的时候很多层提示不支持:
TypeError: ONNX node of type PRelu is not supported.
TypeError: ONNX node of type Shape is not supported
TypeError: ONNX node of type MatMul is not supported.
2、激活函数PRelu修改为Relu之后 重新训练pytorch模型
# 修改激活函数: onnx -->caffe 不支持prelu
self.prelu = ReLU(out_c)
3、type Shape 是在flatten层出现的,修改为torch.reshape还是不行
class Flatten_dfy(Module):
def forward(self, input):
# torch.reshape()
# torch.flatten()
# 返回具有与self相同的数据和元素数量但具有指定形状的张量。
# torch.reshape(input, shape) → Tensor
# return torch.reshape(input, shape=(input.size(0), -1))
return input.view(input.size(0), -1)
后来查资料是要将onnx模型简化:https://github.com/daquexian/onnx-simplifier
一键转换: https://convertmodel.com/ 之后就没有type Shape了
4、type MatMul的 来源是from torch.nn import Linear 就是一个fc层
这个都不支持,实在无语了!最后稍微修改了网络结构,用一个conv2d层(kernel_size=1*1)替换该fc层
注意:保持参数量的不变 之后需要重新训练网络
至此,onnx转caffe终于成功啦!
|
|