【caffe】Caffe模型转换为ONNX模型

在了解了caffe模型的结构和ONNX的结构后,我用python写了一个caffe转onnx的小工具,现只测试了resnet50、alexnet、yolov3的caffe模型和onnx模型推理结果,存在误差,但是在可接受范围内。本工具在转换模型的时候是不需要配置caffe的,只需要安装好protobuf即可。在进行推理测试的时候才需要配置好pycaffe。

github 项目地址

github项目地址
诚恳的求各位博友给个星星O(∩_∩)O~

结构

【caffe】Caffe模型转换为ONNX模型

C2O: 主要代码,caffe_pb2.py为官方caffe.proto编译产生的,caffe_upsample_pb2.py为我为了读取yolov3中upsample层,自行修改caffe.proto文件后编译产生的。Netie用于构建网络结构,Operator用于生成相应的node,createOnnx用于生成onnx模型。
Tool: ModelPath存放了模型路径,test是不带参数的covertONNX,ReadCaffeModel和ReadPrototxt有助于了解caffemodel和prototxt的结构,RunCaffe需要配置pycaffe(对于yolov3,可能还需要相应的cpp、h等文件),RunOnnx用的是OnnxRuntime(这个很容易安装),makeTestData用于生成测试数据。
caffemodel: 存放caffe模型
onnxmodel: 存放onnx模型
proto: 存放proto文件

C2O

主要代码就是C2O中Netie,Operator,createOnnx三个文件

Netie

以Convolution为例(除了一些需要拆分或融合的算子,基本结构都如Conv一样):

1.判断当前layer的类型;

if Layers[i].type == "Convolution" or Layers[i].type == _Layer_CONVOLUTION:

2.获取当前layer的输入名、和输入形状;

inname,input_shape = self.__getLastLayerOutNameAndShape(Layers[i]) #获取输入名列表和输入形状

#获取上一层的输出名(即当前层的输入)
def __getLastLayerOutNameAndShape(self,layer):
    outname = []
    outshape = []
    for i in range(len(layer.bottom)):
        for node in self.NodeList:
            for j in range(len(node.top)):
                if layer.bottom[i] == node.top[j]:
                    name = node.outputs_name[j]
                    shape = node.outputs_shape[j]
        outname.append(name)
        outshape.append(shape)
    return outname,outshape

3.获取当前layer的输出名

outname = self.__getCurrentLayerOutName(Layers[i]) #获取输出名列表

#获取当前层的输出名,即layername+"_Y"
def __getCurrentLayerOutName(self,layer):
    return [layer.name+"_Y"]

4.获取当前layer的输入参数

conv_params = self.__getParams(Layers[i]) #获取输入参数

def __getParams(self, layer):
    for model_layer in self.model_layer:
        if layer.name == model_layer.name:
            Params = copy.deepcopy(model_layer.blobs)
    return Params

5.构建当前layer对应的node,并添加至节点列表

#构建node
conv_node = Operator.createConv(Layers[i], inname, outname, input_shape, conv_params) #构建conv_node
#添加节点到节点列表
self.NodeList.append(conv_node)
self.__n += 1

特殊情况

BatchNorm + Scale —> BatchNormalization
if Layers[i].type == "BatchNorm":
    inname,input_shape = self.__getLastLayerOutNameAndShape(Layers[i])#获取输入名列表和输入形状
    outname = self.__getCurrentLayerOutName(Layers[i]) #获取输出名列表
    bn_params = self.__getParams(Layers[i])[0:2]#获取bn输入参数,由于BatchNorm[2]暂且没用,为了对应BN_paramname,因此只获取[0:2]

    # 构建bn_node
    bn_node = Operator.createBN(Layers[i], inname, outname, input_shape, bn_params)

    scale_params = self.__getParams(Layers[i+1])  # 获取scale输入参数
    scale_node = Operator.createScale(Layers[i + 1], inname, outname, input_shape, scale_params)#构建scale_node

    #将Scale的参数信息添加到BN中,将Scale作为BatchNormalization的主体
    scale_node.params_name.extend(bn_node.params_name)
    scale_node.params_shape.extend(bn_node.params_shape)
    scale_node.params_data.extend(bn_node.params_data)
    #添加节点到节点列表
    self.NodeList.append(scale_node)
    self.__n += 1
InnerProduct —> Reshape + (MatMul + Add) / Gemm
#InnerProduct
# 由于onnx中没有全连接层,因此需要拆分,拆分有两种方法(Reshape+Gemm,Reshape+MatMul+Add)
if Layers[i].type == "InnerProduct" or Layers[i].type == _Layer_INNER_PRODUCT:
    inname,input_shape = self.__getLastLayerOutNameAndShape(Layers[i])#获取reshape的输入名列表和输入形状
    #IP_params = self.__getParams(Layers[i])  # 获取输入参数

    ##reshape
    reshape_layer = copy.deepcopy(Layers[i])#深拷贝
    reshape_layer.name = Layers[i].name+"_Reshape"
    reshape_node = Operator.createReshape(reshape_layer, inname, [reshape_layer.name + "_Y"], input_shape)
    self.NodeList.append(reshape_node)
    self.__n += 1


    # ##gemm未完成
    # gemm_layer = copy.deepcopy(Layers[i])#深拷贝
    # gemm_params = self.__getParams(gemm_layer)
    # gemm_layer.name = Layers[i].name + "_Gemm"
    # #获取input_shape
    # input_shape = self.NodeList[self.__n-1].outputs_shape
    # #构建node
    # gemm_node = C2Omap.createGemm(gemm_layer,[reshape_layer.name+"_Y"],[gemm_layer.name+"_Y"])
    # self.NodeList.append(gemm_node)
    # self.__n += 1


    ##matmul
    matmul_layer = copy.deepcopy(Layers[i])#深拷贝
    matmul_params = self.__getParams(matmul_layer)  # 获取输入参数,对于add来说blobs[1]里存放的是bias不需要,所以直接获取blobs[0]
    matmul_layer.name = Layers[i].name+"_MatMul"
    #获取input_shape
    input_shape = self.NodeList[self.__n-1].outputs_shape
    #构建node
    matmul_node = Operator.createMatmul(matmul_layer, [reshape_layer.name + "_Y"], [matmul_layer.name + "_Y"], input_shape, matmul_params)
    self.NodeList.append(matmul_node)
    self.__n += 1


    ##add
    add_layer = copy.deepcopy(Layers[i])#深拷贝
    add_params = self.__getParams(add_layer) # 获取输入参数,对于add来说blobs[0]里存放的是Weights不需要,所以直接获取blobs[1]
    add_layer.name = Layers[i].name+"_Add"
    #获取input_shape
    input_shape = self.NodeList[self.__n-1].outputs_shape
    #构建node
    add_node = Operator.createAdd(add_layer, [matmul_layer.name + "_Y"], [add_layer.name + "_Y"], input_shape, add_params)
    self.NodeList.append(add_node)
    self.__n += 1

Operator

这个文件里定义了一个c2oNode类,包含了构建onnx node的所有信息,还存储了原有layer的bottom和top(为了方便上面的Netie构建网络结构)

0.c2oNode类

class c2oNode(object):
    def __init__(self,layer,type,inputs_name,outputs_name,inputs_shape,outputs_shape,params_name=[],params_shape = [],params_data=[],dict={}):
        self.__layer = layer
        self.type = type
        self.name = layer.name

        self.bottom = layer.bottom
        self.top = layer.top
        self.inputs_name = inputs_name
        self.outputs_name = outputs_name
        self.inputs_shape = inputs_shape
        self.outputs_shape = outputs_shape

        self.params_name = params_name[0:len(params_shape)]
        self.params_shape = params_shape
        self.params_data = params_data
        self.dict = dict

继续以Conv为例,在Netie调用Operator时,要传入输入名列表、输出名列表、参数[可选]:

1.获取输入参数(不是所有层都会有,如Pooling、Softmax等)

参数名列表、参数维度列表、参数数值列表都要有

#获取输入参数
paramname = [layer.name+"_W",layer.name+"_b"]
paramshape = [p.shape.dim for p in params]
paramdata = [p.data for p in params]

2.获取超参数(也不是所有层都会有,如Eltwise、Reshape等)

超参数的获取,要同时根据onnx所需去找caffe中的param对应

#获取超参数
##填充pads
pads = [0, 0, 0, 0] #默认为0
if layer.convolution_param.pad != []:#若存在pad,则根据pad赋值
    pads = np.array([layer.convolution_param.pad] * 4).reshape(1, -1)[0].tolist()
elif layer.convolution_param.pad_h !=0 or layer.convolution_param.pad_w !=0:#若存在pad_w,pad_h则根据其赋值
    pads = [layer.convolution_param.pad_h,layer.convolution_param.pad_w,layer.convolution_param.pad_h,layer.convolution_param.pad_w]
##步长strides
strides = [1, 1] #默认为1
if layer.convolution_param.stride != []:
    strides = np.array([layer.convolution_param.stride] * 2).reshape(1, -1)[0].tolist()
##卷积核尺寸kernel_shape
kernel_shape = np.array([layer.convolution_param.kernel_size] * 2).reshape(1, -1)[0].tolist()
if layer.convolution_param.kernel_size == []:
    kernel_shape = [layer.convolution_param.kernel_h,layer.convolution_param.kernel_w]
##分组group
group = layer.convolution_param.group
##卷积核数量kernel_num ,用于后面计算输出维度
kernel_num = layer.convolution_param.num_output
#超参数字典
dict = {#"auto_pad":"NOTSET",
        "dilations": [1, 1, 1],
        "group": group,
        "kernel_shape": kernel_shape,
        "pads":pads,
        "strides":strides
        }

3.根据输入维度,计算输出维度

像Conv、Pooling这一类的做输入维度到输出维度的计算,也存在像Softmax、Relu这种输入维度与输出维度相同的layer。

#计算输入维度output_shape
h = (input_shape[0][2] - kernel_shape[0] + 2 * pads[0])/strides[0] + 1 # 输出维度N= ((输入维度I - 卷积核维度K + 2 * 填充P)/步长S) + 1
#当h非整数 ,且未设置pad ,在遇到输出为非整数情况 ,向上取整 ,即在右边和下边补1
if h > int(h) and layer.convolution_param.pad == []:
    output_shape_h = int(h) + 1
    pads = [0,0,1,1]
else:
    output_shape_h = int(h)
output_shape = [[input_shape[0][0],kernel_num,output_shape_h,output_shape_h]]

4.构建layer对应的c2onode

node = c2oNode(layer, "Conv", inname, outname, input_shape, output_shape, paramname, paramshape, paramdata, dict)
其中paramname, paramshape, paramdata, dict根据不同layer类型可选。

createOnnx

1.初始化

def __init__(self):
    self.in_tvi = []#存放输入信息,包括第一个输入和输入参数信息
    self.out_tvi = []#存放输出信息,无特殊情况,只用最后一个输出,但如yolov3这种,有三个输出的特征图
    self.init_t = []#存放输入参数的值
    self.hidden_out_tvi = []#存放中间输出信息

2.创建onnx节点

#创建节点
def __createNode(self, node_type, in_name, out_name, param_name, node_name, dict):
    input_name = copy.deepcopy(in_name)
    if param_name is not None:#如果参数名存在,则将参数名加入输入名列表
        input_name.extend(param_name)
    #print(input_name)
    node_def = helper.make_node(
        node_type,
        input_name,
        out_name,
        node_name,
        **dict,
    )
    return node_def

3.创建模型

#创建模型
def __createModel(self, node_def, graph_name, in_info, out_info, init_info=[],value_info=[]):
    graph_def = helper.make_graph(
        node_def,
        graph_name,
        in_info,
        out_info,
        init_info,
        value_info=value_info
    )
    model_def = helper.make_model(graph_def, producer_name='onnx-example')
    return model_def

4.创建Tensor和TensorValueInfo

def __makeTensorAndValueInfo(self, tensor_name, data, shape,data_type=TensorProto.FLOAT):
    data_tensor = helper.make_tensor(tensor_name, data_type, shape, data)
    data_tensor_value_info = helper.make_tensor_value_info(tensor_name, data_type, shape)
    return data_tensor, data_tensor_value_info

5.获取InputInfo,输入节点信息

def __getInputsInfo(self,firstnode):
    for i in range(len(firstnode.inputs_name)):
        #print(i,firstnode.inputs_name[i])
        input_tvi = helper.make_tensor_value_info(firstnode.inputs_name[i], TensorProto.FLOAT, firstnode.inputs_shape[i])
        self.in_tvi.append(input_tvi)

6.获取OutputInfo,输出节点信息

def __getOutputsInfo(self,lastnode):
    for i in range(len(lastnode.outputs_shape)):
        output_tvi = helper.make_tensor_value_info(lastnode.outputs_name[i],TensorProto.FLOAT,lastnode.outputs_shape[i])
        self.out_tvi.append(output_tvi)

7.获取ValueInfo,中间节点信息

def __getValueInfo(self,innernode):
    for i in range(len(innernode.outputs_shape)):
        hid_out_tvi = helper.make_tensor_value_info(innernode.outputs_name[i], TensorProto.FLOAT, innernode.outputs_shape[i])
        self.hidden_out_tvi.append(hid_out_tvi)

8.获取Initializer,输入参数信息和值

def __getInitializer(self,node):
    for i in range(len(node.params_shape)):
        if node.type == "Reshape":#Reshape的输入参数类型是TensorProto.INT64
            t, tvi = self.__makeTensorAndValueInfo(node.params_name[i], node.params_data[i], node.params_shape[i],TensorProto.INT64)
        else:#其他默认为TensorProto.FLOAT
            t,tvi = self.__makeTensorAndValueInfo(node.params_name[i],node.params_data[i],node.params_shape[i])
        self.init_t.append(t)
        self.in_tvi.append(tvi)

9.判断当前节点是否是输出

def judgeoutput(self,current_node,nodelist):
    for outname in current_node.outputs_name:
        for node in nodelist:
            if outname in node.inputs_name:
                return False
    return True

10.构建节点列表

def __makeNodes(self,HTnodes):
    OnnxNodes = []
    for i in range(len(HTnodes)):
        onnxnode = self.__createNode(HTnodes[i].type,HTnodes[i].inputs_name,HTnodes[i].outputs_name,HTnodes[i].params_name,HTnodes[i].name,HTnodes[i].dict)
        OnnxNodes.append(onnxnode)
        if HTnodes[i].inputs_name[0] == "input":#构建输入节点信息
            self.__getInputsInfo(HTnodes[i])
        elif self.judgeoutput(HTnodes[i],HTnodes):#构建输出节点信息
            self.__getOutputsInfo(HTnodes[i])
        else:#构建中间节点信息
            self.__getValueInfo(HTnodes[i])

        if HTnodes[i].params_name is not None:#如果本节点有参数,那么构建初始化参数
            self.__getInitializer(HTnodes[i])

    return OnnxNodes

11.将caffe转出的节点列表转为onnx

def caffe2onnxmodel(self,nodeslist,modelname):
    node_def = self.__makeNodes(nodeslist)
    print("2:onnx nodes have prepared!")
    onnxmodel = self.__createModel(node_def,modelname,self.in_tvi,self.out_tvi,self.init_t,self.hidden_out_tvi)
    print("3:onnx model conversion successful!")
    return onnxmodel

使用方法

git clone https://github.com/htshinichi/caffe2onnx.git
cd caffe2onnx
python convertONNX.py your_prototxt_path your_caffemodel_path your_onnx_name [--savepath=your_onnx_save_path(optional)]

注意: 若不选择设置onnx_save_path,则需要在Tool/ModelPath中设置root_path

现在支持的算子 Current Support Operator

Convolution
Concat
Dropout
InnerProduct(Reshape+MatMul+Add)
LRN
Pooling
ReLU
Softmax
Eltwise
Upsample
BatchNorm
Scale
PRelu

测试转换用模型 Test Caffe Model

  • Resnet50
  • AlexNet
  • Agenet
  • Yolo V3