"vscode:/vscode.git/clone" did not exist on "d050df368c169e7f94e921fcb3b8b8bec2758176"
Commit 8590ec24 authored by yaoht's avatar yaoht
Browse files

init commit

parent 38a732e1
# Tencent is pleased to support the open source community by making TNN available.
#
# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.
import src.c2oObject as Node
import typing
def getTransposeAttri(layer) -> typing.Dict:
if layer.type == "ShuffleChannel":
# 超参数字典
perm_array = [0, 2, 1, 3, 4]
attributes = {"perm": perm_array}
return attributes
else:
orders = layer.permute_param.order
attributes = {"perm": orders}
return attributes
# 计算输出维度
def getTransposeOutShape(layer, input_shape, attributes):
if layer.type == "ShuffleChannel":
n, g, c, h, w = input_shape[0][0], input_shape[0][1], input_shape[0][2], input_shape[0][3], input_shape[0][4]
output_shape = [[n, c, g, h, w]]
return output_shape
else:
orders = attributes.get("perm")
shape = []
for order in orders:
shape.append(input_shape[0][order])
return [shape]
# 构建节点
def createTranspose(layer, node_name, input_name, output_name, input_shape) -> Node:
attributes = getTransposeAttri(layer)
output_shape = getTransposeOutShape(layer, input_shape, attributes)
node = Node.c2oNode(layer, node_name, "Transpose", input_name, output_name, input_shape, output_shape, attributes)
return node
import numpy as np
import src.c2oObject as Node
##-----------------------------------------------------UnPooling层--------------------------------------------------##
#获取超参数
def getUnPoolingAttri(layer):
# ##池化核尺寸
# kernel_shape = np.array([layer.pooling_param.kernel_size]*2).reshape(1,-1)[0].tolist()
# if layer.pooling_param.kernel_size == []:
# kernel_shape = [layer.pooling_param.kernel_h,layer.pooling_param.kernel_w]
# ##步长
# strides = [1, 1]#默认为1
# if layer.pooling_param.stride != []:
# strides = np.array([layer.pooling_param.stride]*2).reshape(1,-1)[0].tolist()
# ##填充
# pads = [0, 0, 0, 0]#默认为0
# # 这里与卷积时一样,有pad,就按其值设置
# if layer.pooling_param.pad != []:
# pads = np.array([layer.pooling_param.pad] * 4).reshape(1, -1)[0].tolist()
# elif layer.pooling_param.pad_h != 0 or layer.pooling_param.pad_w != 0:
# pads = [layer.pooling_param.pad_h,layer.pooling_param.pad_w,layer.pooling_param.pad_h,layer.pooling_param.pad_w]
#超参数字典
dict = {"kernel_shape": [2, 2],
"strides": [2, 2],
"pads": [0, 0, 0, 0]
}
return dict
#计算输出维度
def getUnPoolingOutShape(input_shape,layer,dict):
kernel_shape = dict["kernel_shape"]
pads = dict["pads"]
strides = dict["strides"]
#计算输出维度,与卷积一样,若为非整数则向上取整
# h = (input_shape[0][2] - kernel_shape[0] + 2 * pads[0])/strides[0] + 1
# if h > int(h):
# output_shape_h = int(h) + 1
# pads = [0,0,1,1]
# else:
# output_shape_h = int(h)
# output_shape = [[input_shape[0][0],input_shape[0][1],output_shape_h,output_shape_h]]
output_shape = [[input_shape[0][0], input_shape[0][1], input_shape[0][2]*2, input_shape[0][3]*2]]
return output_shape
#构建节点
def createUnPooling(layer,nodename,inname,outname,input_shape):
dict = getUnPoolingAttri(layer)
output_shape = getUnPoolingOutShape(input_shape,layer,dict)
node = Node.c2oNode(layer, nodename, "MaxUnpool", inname, outname, input_shape, output_shape, dict=dict)
return node
import src.c2oObject as Node
import numpy as np
# 获取超参数
def get_upsample_attri(layer):
# scale = layer.upsample_param.scale
# scales = [1.0,1.0,scale,scale]
# dict = {"scales":scales,"mode":"nearest"}#Upsample将scales放入参数里面了
# dict = {"width_scale": scale,"height_scale":scale, "mode": "nearest"}#在OpenVINO读onnx的时候要求用width_scale和height_scale
scale = layer.upsample_param.scale
scales = [1.0, 1.0, scale, scale]
attributes = {"mode": "linear",
'scales': scales}
return attributes
def get_upsample_outputshape(input_shape, layer):
scale = layer.upsample_param.scale
scales = [1.0, 1.0, scale, scale]
output_shape = [np.multiply(np.array(scales, dtype=np.int), np.array(input_shape[0])).tolist()]
return output_shape
def create_upsample_node(layer, node_name, input_name, output_name, input_shape):
attributes = get_upsample_attri(layer)
output_shape = get_upsample_outputshape(input_shape, layer)
# print(output_shape)
node = Node.c2oNode(layer, node_name, "Upsample", input_name, output_name, input_shape, output_shape, attributes)
return node
from src.OPs.BatchNorm import *
from src.OPs.Concat import *
from src.OPs.Conv import *
from src.OPs.Dropout import *
from src.OPs.Eltwise import *
from src.OPs.Gemm import *
from src.OPs.LRN import *
from src.OPs.Pooling import *
from src.OPs.PRelu import *
from src.OPs.ReLU import *
from src.OPs.Reshape import *
from src.OPs.Softmax import *
from src.OPs.Upsample import *
from src.OPs.UnPooling import *
from src.OPs.ConvTranspose import *
from src.OPs.Slice import *
from src.OPs.Transpose import *
from src.OPs.Sigmoid import *
from src.OPs.Min import *
from src.OPs.Clip import *
from src.OPs.Log import *
from src.OPs.Mul import *
from src.OPs.Interp import *
from src.OPs.Crop import *
from src.OPs.InstanceNorm import *
from src.OPs.PriroBox import create_priorbox_node
from src.OPs.DetectionOutput import create_detection_output
from src.OPs.Flatten import create_flatten_node
from src.OPs.Resize import create_resize_node
from src.OPs.Axpy import create_axpy_add_node, create_axpy_mul_node
from src.OPs.LpNormalization import create_Lp_Normalization
from src.OPs.Power import get_power_param, create_power_node
from src.OPs.Add import create_add_node
from src.OPs.Tanh import createTanh
# Tencent is pleased to support the open source community by making TNN available.
#
# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.
import argparse
def parse_args():
parser = argparse.ArgumentParser(description='convert caffe model to onnx')
parser.add_argument(dest='proto_file',
action='store',
help='the path for prototxt file, the file name must end with .prototxt')
parser.add_argument(dest='caffe_model_file',
action='store',
help='the path for caffe model file, the file name must end with .caffemodel!')
parser.add_argument('-o',
dest='onnx_file',
action='store',
help='the path for generate onnx file')
args = parser.parse_args()
return args
from onnx import helper
class c2oNode(object):
def __init__(self,layer,node_name,type,inputs_name,outputs_name,inputs_shape,outputs_shape,dict={}, Flag=False):
self.node = self.__createNode(type, inputs_name, outputs_name, node_name, dict)
self.bottom = layer.bottom
if Flag is True:
self.top = outputs_name
else:
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.dict = dict
#创建节点
def __createNode(self, node_type, in_name, out_name, node_name, dict):
node_def = helper.make_node(
node_type,
in_name,
out_name,
node_name,
**dict,
)
return node_def
class c2oGraph():
def __init__(self,onnxname):
self.name = onnxname
self.in_tvi = []#存放输入信息,包括第一个输入和输入参数信息
self.out_tvi = []#存放输出信息
self.init_t = []#存放输入参数的值
self.hidden_out_tvi = []#存放中间输出信息
def addInputsTVI(self,in_tvi):
self.in_tvi.append(in_tvi)
def addOutputsTVI(self,out_tvi):
self.out_tvi.append(out_tvi)
def addInitTensor(self,init_t):
self.init_t.append(init_t)
def addValueInfoTVI(self,vi_tvi):
self.hidden_out_tvi.append(vi_tvi)
import src.OPs as op
from src.c2oObject import *
from onnx import helper
import copy
import numpy as np
from src.op_layer_info import *
import random
import sys
from typing import *
import onnx
class Caffe2Onnx():
def __init__(self, net, model, onnxname):
# 初始化一个c2oGraph对象
self.onnxmodel = c2oGraph(onnxname)
# 网络和参数
self.netLayerCaffe = self.GetNetLayerCaffe(net)
self.netModelCaffe = self.GetNetModelCaffe(model)
# 模型的输入名和输入维度
self.model_input_name = []
self.model_input_shape = []
# 节点列表
self.onnxNodeList = []
# 获取层列表
LayerList = self.AddInputsTVIAndGetLayerList(net)
self.GenerateOnnxNodeList(LayerList)
self.AddOutputsTVIAndValueInfo()
# 获取网络层
def GetNetLayerCaffe(self, net):
if len(net.layer) == 0 and len(net.layers) != 0:
return net.layers
elif len(net.layer) != 0 and len(net.layers) == 0:
return net.layer
else:
print("prototxt layer error")
return -1
# 获取参数层
def GetNetModelCaffe(self, model):
if len(model.layer) == 0 and len(model.layers) != 0:
return model.layers
elif len(model.layer) != 0 and len(model.layers) == 0:
return model.layer
else:
print("caffemodel layer error")
return -1
# 将模型输入信息添加到Inputs中并获取后续层列表
def AddInputsTVIAndGetLayerList(self, net):
# 如果第一个layer的类型为Input,且没有net.input存在
if net.input == [] and self.netLayerCaffe[0].type == "Input":
layer_list = []
# 考虑到整个网络会有多输入情况
for lay in self.netLayerCaffe:
if lay.type == "Input":
if len(lay.top) == 1 and lay.top[0] != lay.name:
input_layer_name = lay.top[0]
else:
input_layer_name = lay.name
in_tvi = helper.make_tensor_value_info(
input_layer_name + "_input", TensorProto.FLOAT,
lay.input_param.shape[0].dim)
self.model_input_name.append(input_layer_name + "_input")
self.model_input_shape.append(lay.input_param.shape[0].dim)
self.onnxmodel.addInputsTVI(in_tvi)
else:
layer_list.append(lay)
return layer_list
# 如果存在net.input
elif net.input != []:
if bool(net.input_dim):
input_dim = net.input_dim
elif bool(net.input_shape):
input_dim = net.input_shape[0].dim
else:
raise RuntimeError("Input shape missing!")
in_tvi = helper.make_tensor_value_info("input", TensorProto.FLOAT, input_dim)
self.model_input_name.append("input")
self.model_input_shape.append(input_dim)
self.onnxmodel.addInputsTVI(in_tvi)
return self.netLayerCaffe
# 以上情况都不是,则该caffe模型没有输入,存在问题
else:
raise ValueError("the caffe model has no input")
# 得到layer的参数shape
def GetParamsShapeAndData(self, layer):
ParamShape = []
ParamData = []
# 根据这个layer名找出对应的caffemodel中的参数
for model_layer in self.netModelCaffe:
if layer.name == model_layer.name:
Params = copy.deepcopy(model_layer.blobs)
ParamShape = [p.shape.dim for p in Params]
ParamData = [p.data for p in Params]
if layer.type == "BatchNorm" or layer.type == "BN":
if len(ParamShape) == 3:
# 如果是bn层,则不用最后一层的滑动系数
ParamShape = ParamShape[:-1]
ParamData = ParamData[:-1]
elif len(ParamShape) == 2 and len(ParamShape[0]) != 1:
ParamShape = [[ParamShape[0][1]], [ParamShape[1][1]]]
ParamData = ParamData
return ParamShape, ParamData
def get_param_shape(self, params):
shapes = []
for p in params:
if p.shape.dim != []:
shape = p.shape.dim
shapes.append(shape)
else:
shape = [p.num, p.channels, p.height, p.width]
shapes.append(shape)
return shapes
# 将参数添加到Inputs中,并生成tensor存储数据
def AddInputsTVIFromParams(self, layer, ParamName, ParamType):
ParamShape = []
ParamData = []
# 根据这个layer名找出对应的caffemodel中的参数
for model_layer in self.netModelCaffe:
if layer.name == model_layer.name:
Params = copy.deepcopy(model_layer.blobs)
#ParamShape = [p.shape.dim for p in Params]
ParamShape = self.get_param_shape(Params)
ParamData = [p.data for p in Params]
if layer.type == "BatchNorm" or layer.type == "BN":
if len(ParamShape) == 3:
# 如果是bn层,params为[mean, var, s],则需要把mean和var除以滑动系数s
ParamShape = ParamShape[:-1]
ParamData = [
[q / (Params[-1].data[0])
for q in p.data] if i == 0 else
[q / (Params[-1].data[0] + 1e-5) for q in p.data]
for i, p in enumerate(Params[:-1])
] # with s
elif len(ParamShape) == 2 and len(ParamShape[0]) == 4:
ParamShape = [[ParamShape[0][1]], [ParamShape[1][1]]]
ParamData = [[q / 1. for q in p.data] if i == 0 else
[q / (1. + 1e-5) for q in p.data]
for i, p in enumerate(Params)]
if layer.type == "Reshape":
ParamShape = [[len(model_layer.reshape_param.shape.dim)]]
ParamData = [model_layer.reshape_param.shape.dim]
if layer.type == "Convolution" or layer.type == "ConvolutionDepthwise":
if len(ParamShape) == 2:
ParamShape[1] = [ParamShape[0][0]]
if layer.type == "InnerProduct":
if len(ParamShape[0]) > 2:
ParamShape[0] = [ParamShape[0][2], ParamShape[0][3]]
if len(ParamShape) == 2:
if len(ParamShape[1]) > 2:
ParamShape[1] = [ParamShape[1][2], ParamShape[1][3]]
if layer.type == "Normalize":
if len(ParamShape) == 1:
ParamShape[0] = [1, ParamShape[0][0], 1, 1]
# comment it for tvm because tvm use broadcast at prelu layer
# 个人感觉如果不用 tvm,就不需要使用 Prelu
# if layer.type == 'PReLU':
# ParamShape = [[ParamShape[0][0], 1, 1]]
break
# 判断是否有Param
if ParamShape != []:
ParamName = ParamName[0:len(ParamShape)]
ParamType = ParamType[0:len(ParamShape)]
for i in range(len(ParamShape)):
ParamName[i] = layer.name + ParamName[i]
p_tvi = helper.make_tensor_value_info(ParamName[i],
ParamType[i],
ParamShape[i])
p_t = helper.make_tensor(ParamName[i], ParamType[i],
ParamShape[i], ParamData[i])
self.onnxmodel.addInputsTVI(p_tvi)
self.onnxmodel.addInitTensor(p_t)
# print("添加参数" + ParamName[i] + "输入信息和tensor数据")
if layer.type == "BatchNorm" or layer.type == "BN" or layer.type == "Scale":
return ParamName, ParamShape
return ParamName
# 手动将参数添加到输入信息中,并生成tensor存储数据
def AddInputsTVIMannul(self, layer, param_names, param_types, param_shapes,
param_data):
node_names = copy.deepcopy(param_names)
for i in range(len(param_shapes)):
node_names[i] = layer.name + param_names[i]
p_tvi = helper.make_tensor_value_info(node_names[i],
param_types[i],
param_shapes[i])
p_t = helper.make_tensor(node_names[i], param_types[i],
param_shapes[i], param_data[i])
self.onnxmodel.addInputsTVI(p_tvi)
self.onnxmodel.addInitTensor(p_t)
return node_names
# # 由于 Slice 的 input 情况特殊,所以需要特殊处理
# if layer.type == 'Slice':
# for i in range(len(ParamShape)):
# p_tvi = helper.make_tensor_value_info(Param_Name[i], ParamType[i], ParamShape[i])
# p_t = helper.make_tensor(Param_Name[i], ParamType[i], ParamShape[i], ParamData[i])
# self.onnxmodel.addInputsTVI(p_tvi)
# self.onnxmodel.addInitTensor(p_t)
# return Param_Name
# else:
# for i in range(len(ParamShape)):
# Param_Name[i] = layer.name + ParamName[i]
# p_tvi = helper.make_tensor_value_info(Param_Name[i], ParamType[i], ParamShape[i])
# p_t = helper.make_tensor(Param_Name[i], ParamType[i], ParamShape[i], ParamData[i])
# self.onnxmodel.addInputsTVI(p_tvi)
# self.onnxmodel.addInitTensor(p_t)
# return Param_Name
# 获取上一层的输出名(即当前层的输入)
def GetLastLayerOutNameAndShape(self, layer):
output_name = []
outshape = []
# flag is True: 模型的输入没有被覆盖
# flag is False: 模型的输入已经被覆盖
flag = True
# 如果结点列表为空,或者当前层的bottom在input_name中,那么上一层输入一定是 Input
if self.onnxNodeList == []:
output_name += self.model_input_name
outshape += self.model_input_shape
else:
for i in range(len(layer.bottom)):
# 因为prototxt中存在top和bottom同名的情况,但是layer.bottom只能对应一个node,所以对每个layer.bottom,找到最末的那个同名节点作为上一层节点
name = None
shape = None
for node in self.onnxNodeList:
for j in range(len(node.top) if node.node.op_type != "MaxPool" else 1):
if layer.bottom[i] == node.top[j]:
name = node.outputs_name[j]
shape = node.outputs_shape[j]
for k in range(len(node.bottom)):
if node.top[j] == node.bottom[k]:
for w in range(len(self.model_input_name)):
if node.top[j] + '_input' == self.model_input_name[w]:
flag = False
for j in range(len(self.model_input_name)):
if layer.bottom[i] + '_input' == self.model_input_name[j] and flag:
output_name.append(self.model_input_name[j])
outshape.append(self.model_input_shape[j])
if name:
output_name.append(name)
outshape.append(shape)
try:
assert output_name, "Failed at layer %s, layer's bottom not detected ..." % (layer.name)
except:
print("Failed at layer %s, layer's bottom not detected ..." % (layer.name))
exit(-1)
return output_name, outshape
# 获取当前层的输出名,即layername
def GetCurrentLayerOutName(self, layer):
# return [layer.name]
# 考虑有多个输出的情况
# # TODO: 为什么要使用 layer.name 进行替代呢?
if layer.top == layer.bottom and len(layer.top) == 1:
return [layer.name]
return [out for out in layer.top]
def GenerateOnnxNodeList(self, Layers):
for i in range(len(Layers)):
print("convert layer: " + Layers[i].name)
# Convolution
if Layers[i].type == "Convolution" or Layers[i]. type == Layer_CONVOLUTION:
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
conv_pname = self.AddInputsTVIFromParams(Layers[i], op_pname["Conv"], op_ptype["Conv"])
input_name.extend(conv_pname)
# 3.构建conv_node
conv_node = op.createConv(Layers[i], node_name, input_name, output_name, input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(conv_node)
elif Layers[i].type == "ConvolutionDepthwise" or Layers[i].type == Layer_CONVOLUTION:
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
conv_pname = self.AddInputsTVIFromParams(Layers[i], op_pname["Conv"], op_ptype["Conv"])
input_name.extend(conv_pname)
# 3.构建conv_node
conv_node = op.createConv(Layers[i], node_name, input_name, output_name, input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(conv_node)
# BatchNorm+Scale
elif Layers[i].type == "BatchNorm" or Layers[i].type == "BN":
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
node_name = Layers[i].name
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
if i < len(Layers) - 1 and Layers[i + 1].type == "Scale":
scale_pname, scale_pshape = self.AddInputsTVIFromParams(Layers[i + 1], op_pname["Scale"],
op_ptype["Scale"])
bn_pname, bn_pshape = self.AddInputsTVIFromParams(Layers[i], op_pname["BatchNorm"],
op_ptype["BatchNorm"])
assert bn_pshape == scale_pshape, "BatchNorm and Scale params should share the same shape"
input_name.extend(scale_pname)
input_name.extend(bn_pname)
else:
bn_pshape, _ = self.GetParamsShapeAndData(Layers[i])
custom_params = [np.ones(shape=bn_pshape[0], dtype=np.float),
0.001 + np.zeros(shape=bn_pshape[1], dtype=np.float)]
scale_pname = self.AddInputsTVIMannul(Layers[i], op_pname["Scale"], op_ptype["Scale"], bn_pshape,
custom_params)
bn_pname, bn_pshape = self.AddInputsTVIFromParams(Layers[i], op_pname["BatchNorm"],
op_ptype["BatchNorm"])
input_name.extend(scale_pname)
input_name.extend(bn_pname)
# 3.构建bn_node
bn_node = op.createBN(Layers[i], node_name, input_name, output_name, input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(bn_node)
elif Layers[i].type == "Scale":
if i > 0 and (Layers[i - 1].type == "BatchNorm" or Layers[i - 1].type == "BN"):
# bn + scale
continue
# signal scale
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
# node_name = Layers[i].name + random.choice('1234567890abcdefghijklmnopqrst')
node_name = Layers[i].name
has_two_input: bool = False
if len(input_name) > 1:
has_two_input = True
if has_two_input and op.need_add_reshape(input_shape):
reshape_layer = copy.deepcopy(Layers[i])
# add reshape layer
reshape_node_name = input_name[1] + '_reshap_' + random.choice('1234567890abcdefghijklmnopqrst')
reshape_input_name = input_name[1]
reshape_input_shape = input_shape[1]
reshape_shape_data = op.get_param_shape(input_shape)
reshape_shape_shape = np.shape(reshape_shape_data)
reshape_params = self.AddInputsTVIMannul(Layers[i], [reshape_node_name + 'shape'], [TensorProto.INT64],
[reshape_shape_shape], [reshape_shape_data])
reshape_output_name = [reshape_input_name + '_output_name']
reshape_node = op.createReshape(reshape_layer, reshape_node_name, [reshape_input_name, reshape_params[0]],
reshape_output_name, reshape_input_shape, output_shape=[reshape_shape_data])
self.onnxNodeList.append(reshape_node)
# add mul node
input_name[1] = reshape_output_name[0]
input_shape[1] = reshape_shape_data
mul_node = op.create_mul_node(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(mul_node)
else:
param_shape, param_data = self.GetParamsShapeAndData(Layers[i])
# Scale = Mul + Add
if len(param_shape) == 2:
# create mul
param_scale_shape = [1, param_shape[0][0], 1, 1]
param_scale_data = param_data[0]
param_scale_name = self.AddInputsTVIMannul(Layers[i], ["_scale"], [TensorProto.FLOAT], [param_scale_shape], [param_scale_data])
mul_node_name = node_name + "_mul"
mul_input_name = [input_name[0], param_scale_name[0]]
mul_output_name = [output_name[0] + "_mul"]
mul_input_shape = [input_shape[0], param_scale_shape]
mul_node = op.create_mul_node(Layers[i], mul_node_name, mul_input_name, mul_output_name, mul_input_shape)
self.onnxNodeList.append(mul_node)
param_bias_shape = [1, param_shape[1][0], 1, 1]
param_bias_data = param_data[1]
param_bias_name = self.AddInputsTVIMannul(Layers[i], ["_bias"], [TensorProto.FLOAT], [param_bias_shape], [param_bias_data])
add_node_name = node_name + "_add"
add_input_name = [mul_output_name[0], param_bias_name[0]]
add_output_name = output_name
add_input_shape = [input_shape[0], param_bias_shape]
add_node = op.create_add_node(Layers[i], add_node_name, add_input_name, add_output_name, add_input_shape)
self.onnxNodeList.append(add_node)
# Scale = Mul
if len(param_shape) == 1:
# create mul
param_scale_shape = [1, param_shape[0][0], 1, 1]
param_scale_data = param_data[0]
param_scale_name = self.AddInputsTVIMannul(
Layers[i], ["_scale"], [TensorProto.FLOAT],
[param_scale_shape], [param_scale_data])
mul_input_name = [input_name[0], param_scale_name[0]]
mul_input_shape = [input_shape[0], param_scale_shape]
mul_node = op.create_mul_node(Layers[i], node_name,
mul_input_name,
output_name,
mul_input_shape)
self.onnxNodeList.append(mul_node)
# Pooling
elif Layers[i].type == "Pooling" or Layers[i].type == Layer_POOLING:
# TODO:
# Pooling <= Pad + Pool
# NOTE: 由于 Caffe 和 ONNX 对 AveragePool 的处理的方式的不同,所以需要在pool node 之前添加 Pad node
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
node_name = Layers[i].name
# create pad node
pads = op.get_pool_pads(Layers[i])
pads_shape = [np.shape(pads)]
pads_name = node_name + "_output"
pads_output_name = [node_name + "_output"]
pad_output_shape = op.calculate_pad_output_shape(input_shape, pads)
pads_param = self.AddInputsTVIMannul(Layers[i], ["_pad"], [TensorProto.INT64], pads_shape, [pads])
input_name.extend(pads_param)
pool_type = op.pooling_type(Layers[i])
if pool_type == "GlobalMaxPool" or pool_type == "MaxPool":
constant_value = [-sys.float_info.max]
constant_shape = [np.shape(constant_value)]
constant_value_param = self.AddInputsTVIMannul(Layers[i], ["_constant_value"], [TensorProto.FLOAT],
constant_shape, [constant_value])
input_name.extend(constant_value_param)
pad_node = op.create_pad_node(Layers[i], pads_name, input_name, pads_output_name, input_shape)
self.onnxNodeList.append(pad_node)
# 2.构建pool_node
pool_node = op.create_pooling_node(Layers[i], node_name, pads_output_name, output_name,
pad_output_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(pool_node)
# MaxUnPool
elif Layers[i].type == "MaxUnpool":
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
node_name = Layers[i].name
# 2.构建unpool_node
unpool_node = op.createUnPooling(Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(unpool_node)
# Eltwise
elif Layers[i].type == "Eltwise" or Layers[i].type == Layer_ELTWISE:
# 1.获取节点输入名、输入维度、输出名、节点名
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
node_name = Layers[i].name
# 2.构建eltwise_node
eltwise_node = op.createEltwise(Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(eltwise_node)
# Softmax
elif Layers[i].type == "Softmax" or Layers[i].type == Layer_SOFTMAX:
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
node_name = Layers[i].name
# 2.构建softmax_node
softmax_node = op.createSoftmax(Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(softmax_node)
# Relu
elif Layers[i].type == "ReLU" or Layers[i].type == Layer_RELU:
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
node_name = Layers[i].name
# letters = '1234567890abcdefghijklmnopqrst'
# length = random.randrange(5, 16)
# randstr = ''.join(random.choice(letters) for _ in range(length))
# node_name = node_name
# for i in range(len(output_name)):
# output_name[i] = output_name[i] + random.choice('1234567890abcdef')
#print(output_name)
# 2.构建relu_node
relu_node = op.createRelu(Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(relu_node)
# PRelu
elif Layers[i].type == "PReLU":
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
pname = self.AddInputsTVIFromParams(Layers[i], op_pname["PRelu"], op_ptype["PRelu"])
input_name.extend(pname)
# 3.构建PRelu_node
PRelu_node = op.createPRelu(Layers[i], node_name, input_name, output_name, input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(PRelu_node)
# relu6
elif Layers[i].type == 'ReLU6':
# relu6 = clip(0, 6)
# add relu node
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
min_value = np.float(0)
max_value = np.float(6)
shape = np.shape([min_value])
min_param = self.AddInputsTVIMannul(Layers[i], ["_min"],
[TensorProto.FLOAT], [shape],
[[min_value]])
input_name.extend(min_param)
max_param = self.AddInputsTVIMannul(Layers[i], ['_max'],
[TensorProto.FLOAT], [shape],
[[max_value]])
input_name.extend(max_param)
relu6_node = op.create_clip_node(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(relu6_node)
elif Layers[i].type == "Sigmoid":
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
node_name = Layers[i].name
# 2.构建relu_node
sigmoid_node = op.createSigmoid(Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(sigmoid_node)
elif Layers[i].type == 'Log':
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
node_name = Layers[i].name
log_node = op.create_log_node(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(log_node)
# LRN
elif Layers[i].type == "LRN" or Layers[i].type == Layer_LRN:
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.构建LRN_node
LRN_node = op.createLRN(Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(LRN_node)
# Dropout
elif Layers[i].type == "Dropout" or Layers[i].type == Layer_DROPOUT:
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.构建Dropout_node
Dropout_node = op.createDropout(Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(Dropout_node)
# Upsample
elif Layers[i].type == "Upsample" or Layers[i].type == Layer_UPSAMPLE:
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
# add roi input
# add scales input
paramshape = [[8, 1],
[4, 1]]
paramdata = [[1, 1, 1, 1, 2, 2, 2, 2],
[1.0, 1.0, Layers[i].upsample_param.scale, Layers[i].upsample_param.scale]]
pname = self.AddInputsTVIMannul(Layers[i], op_pname["Upsample"], op_ptype["Upsample"], paramshape,
paramdata)
input_name.extend(pname)
# 3.构建Upsample_node
Upsample_node = op.create_resize_node(Layers[i], node_name, input_name, output_name, input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(Upsample_node)
elif Layers[i].type == 'Interp':
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
interp_node = op.create_interp_node(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(interp_node)
# Concat
elif Layers[i].type == "Concat" or Layers[i].type == Layer_CONCAT:
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.构建Concat_node
Concat_node = op.createConcat(Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(Concat_node)
elif Layers[i].type == 'Slice':
# 1. 获取节点书输入名,输入维度,输出名,节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name_list = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
starts, ends, axes = op.analyzeLayer(Layers[i], input_shape)
SliceLayer = copy.deepcopy(Layers[i])
for i in range(len(output_name_list)):
# 放在这里的原因是
slice_name = copy.deepcopy(input_name)
# starts ends axes 的 shape 是相同的
shape = [np.shape([1])]
starts_param = self.AddInputsTVIMannul(SliceLayer, ['_starts' + str(i)],
[TensorProto.INT64], shape,
[[starts[i]]])
ends_param = self.AddInputsTVIMannul(SliceLayer, ['_ends' + str(i)],
[TensorProto.INT64], shape,
[[ends[i]]])
axes_param = self.AddInputsTVIMannul(SliceLayer, ['_axes' + str(i)],
[TensorProto.INT64], shape,
[[axes[i]]])
slice_name.extend(starts_param)
slice_name.extend(ends_param)
slice_name.extend(axes_param)
Slice_node = op.createSlice(SliceLayer, output_name_list[i], slice_name, [output_name_list[i]],
input_shape, starts[i], ends[i])
# 3. 添加节点到节点列表
self.onnxNodeList.append(Slice_node)
# Reshape
elif Layers[i].type == "Reshape":
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
reshape_param = op.get_reshape_param(Layers[i], input_shape)
reshape_param_shape = [np.shape(reshape_param)]
pname = self.AddInputsTVIMannul(Layers[i], op_pname["Reshape"], op_ptype["Reshape"], reshape_param_shape,
[reshape_param])
input_name.extend(pname)
# 3.构建reshape节点
reshape_node = op.createReshape(Layers[i], node_name, input_name, output_name, input_shape)
# 4.添加点到节点列表
self.onnxNodeList.append(reshape_node)
# InnerProduct
# 由于onnx中没有全连接层,因此需要拆分,拆分有两种方法(Reshape+Gemm,Reshape+MatMul+Add)
elif Layers[i].type == "InnerProduct" or Layers[i].type == Layer_INNER_PRODUCT:
node_layer = copy.deepcopy(Layers[i]) # 深拷贝
node_input_name, node_input_shape = self.GetLastLayerOutNameAndShape(node_layer) # 获取输入名列表和输入形状
reshape_outname = ""
reshape_output_shape = op.getReshapeOutShape(Layers[i], node_input_shape)
need_reshape = 0 if reshape_output_shape[0] == node_input_shape[0] else 1
if need_reshape:
####一、reshape
# 1.获取节点输入名、输入维度、输出名、节点名
reshape_outname = [node_layer.name + "_Reshape"]
reshape_nodename = node_layer.name + "_Reshape"
# 2.生成节点参数tensor value info,并获取节点参数名, 将参数名加入节点输入名列表
paramshape = [[2]]
reshape_pname = self.AddInputsTVIMannul(node_layer, op_pname["Reshape"], op_ptype["Reshape"],
paramshape, reshape_output_shape)
node_input_name.extend(reshape_pname)
# 3.构建reshape_node
reshape_node = op.createReshape(node_layer, reshape_nodename, node_input_name, reshape_outname,
node_input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(reshape_node)
# import ipdb; ipdb.set_trace()
####二、Gemm 最后一个node输出保持原名称
gemm_layer = copy.deepcopy(Layers[i]) # 深拷贝
# 1.获取节点输入名、输入维度、输出名、节点名
gemm_inname = reshape_outname if need_reshape == 1 else node_input_name
gemm_input_shape = reshape_output_shape if need_reshape == 1 else node_input_shape
gemm_outname = [gemm_layer.name]
gemm_nodename = gemm_layer.name
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
gemm_pname = self.AddInputsTVIFromParams(gemm_layer, op_pname["InnerProduct"], op_ptype[
"InnerProduct"]) # 获取输入参数,对于add来说blobs[1]里存放的是bias不需要,所以直接获取blobs[0]
gemm_inname.extend(gemm_pname)
# 3.构建gemm_node
matmul_node = op.createGemm(gemm_layer, gemm_nodename, gemm_inname, gemm_outname, gemm_input_shape,
gemm_layer.inner_product_param.num_output)
# 4.添加节点到节点列表
self.onnxNodeList.append(matmul_node)
elif Layers[i].type == 'ShuffleChannel':
# TODO support ShuffleChannel
# reshape [N, C, H, W] tensor to [N, G, C', H, W]
node_layer = copy.deepcopy(Layers[i]) # 深拷贝
node_input_name, node_input_shape = self.GetLastLayerOutNameAndShape(node_layer) # 获取输入名列表和输入形状
reshape_outname = ""
reshape_output_shape = op.getReshapeOutShape(Layers[i], node_input_shape)
need_reshape = 0 if reshape_output_shape[0] == node_input_shape[0] else 1
if need_reshape:
# 一. reshape [N, C, H, W] tensor to [N, G, C', H, W]
# 1.获取节点输入名、输入维度、输出名、节点名
reshape_outname = [node_layer.name + "_Reshape"]
reshape_nodename = node_layer.name + "_Reshape"
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
param_data = op.getReshapeOutShape(node_layer, node_input_shape)
param_shape = np.array([1, 2, 3, 4, 5], np.int).shape
reshape_pname = self.AddInputsTVIMannul(node_layer, op_pname["Reshape"], op_ptype["Reshape"],
[param_shape], param_data)
node_input_name.extend(reshape_pname)
# 这里不用对输入进行拓展,因为输入没有增加
# node_input_name.extend(reshape_pname)
# 3.构建reshape_node
reshape_node = op.createReshape(node_layer,
reshape_nodename,
node_input_name,
reshape_outname,
node_input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(reshape_node)
# 2. transpose [N, C', G, H, W]
transpose_layer = copy.deepcopy(Layers[i]) # 深拷贝
# 1.获取节点输入名、输入维度、输出名、节点名
transpose_input_name = reshape_outname if need_reshape == 1 else node_input_name
transpose_input_shape = reshape_output_shape if need_reshape == 1 else node_input_shape
transpose_output_name = [node_layer.name + "_Transpose"]
transpose_node_name = node_layer.name + "_Transpose"
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
# 获取输入参数,对于add来说blobs[1]里存放的是bias不需要,所以直接获取blobs[0]
# TODO 这地方为什么要选择使用AddInputsTVIMannul?取决于什么?
# ANSWER: 取决于要转换的 onnx 的类型
# TODO param_date 是什么?为什么要设置这个变量
param_data = [[2]]
# transpose_pname = self.AddInputsTVIMannul(transpose_layer,
# op_pname["Transpose"],
# op_ptype['Transpose'],
# param_data,
# transpose_input_shape)
# transpose_input_name.extend(transpose_pname)
# 3.
transpose_node = op.createTranspose(transpose_layer,
transpose_node_name,
transpose_input_name,
transpose_output_name,
transpose_input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(transpose_node)
# 三、 Reshape [N, C', G, H, W] tensor to [N, C, H, W]
#
end_layer = copy.deepcopy(Layers[i])
end_layer.type = "DeReshape"
# 最后的输出的节点要保持原名称,这是为了生成该节点,保持链路畅通
end_output_name = [end_layer.name]
end_node_name = end_layer.name
# 上一层的输出是这一层的输入
end_input_name = transpose_node.outputs_name
end_input_shape = transpose_node.outputs_shape
# 最后保持输出和输入的形状是一致的
end_output_shape = [[node_input_shape[0][0], -1, node_input_shape[0][2], node_input_shape[0][3]]]
param_shape = [np.array([1, 2, 3, 4], dtype=np.int).shape]
end_pname = self.AddInputsTVIMannul(node_layer, op_pname["DouReshape"], op_ptype["DouReshape"],
param_shape, end_output_shape)
end_input_name.extend(end_pname)
# 构建
end_node = op.createReshape(end_layer,
end_node_name,
end_input_name,
end_output_name,
end_input_shape)
self.onnxNodeList.append(end_node)
# Deconvolution
elif Layers[i].type == "Deconvolution":
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
conv_pname = self.AddInputsTVIFromParams(Layers[i], op_pname["ConvTranspose"],
op_ptype["ConvTranspose"])
input_name.extend(conv_pname)
# 3.构建conv_node
conv_node = op.createConvTranspose(Layers[i], node_name, input_name, output_name, input_shape)
# if True:
# self.__print_debug_info(node_name, input_name, output_name, input_shape, conv_node.outputs_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(conv_node)
# Flatten
elif Layers[i].type == "Flatten":
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# 由于后面 Flatten 的优化有问题,所以目前先将 Flatten -> reshape
# flatten_node = op.create_flatten_node(layers[i], node_name, input_name,
# output_name, input_shape)
# self.onnxnodelist.append(flatten_nodelatten_node)
# continue
# Flatten -> Reshape
# import ipdb; ipdb.set_trace()
# # 2.生成节点参数tensor value info,并获取节点参数名,将参数名加入节点输入名列表
paramshape = [[2]]
paramdata = op.getReshapeOutShape(Layers[i], input_shape)
reshape_pname = self.AddInputsTVIMannul(Layers[i], op_pname["Reshape"], op_ptype["Reshape"], paramshape,
paramdata)
input_name.extend(reshape_pname)
# 3.构建reshape_node
reshape_node = op.createReshape(Layers[i], node_name, input_name, output_name, input_shape)
# 4.添加节点到节点列表
self.onnxNodeList.append(reshape_node)
elif Layers[i].type == "Permute":
# Permute -> Transpose
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
transpose_node = op.createTranspose(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(transpose_node)
elif Layers[i].type == "PriorBox":
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
priorbox_node = op.create_priorbox_node(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(priorbox_node)
elif Layers[i].type == "DetectionOutput":
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
detection_output_node = op.create_detection_output(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(detection_output_node)
elif Layers[i].type == "Axpy":
# axpy = mul + add
# top = bottom[0] * bottom[1] + bottom[2]
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
# create mul node
mul_node = op.create_axpy_mul_node(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(mul_node)
# create add node
add_node = op.create_axpy_add_node(Layers[i], node_name, input_name, output_name, input_shape)
self.onnxNodeList.append(add_node)
elif Layers[i].type == "Normalize":
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
lp_normalization_output_name = [output_name[0] + "_lp"]
lp_normalization_node = op.create_Lp_Normalization(Layers[i], node_name, input_name,
lp_normalization_output_name, input_shape)
self.onnxNodeList.append(lp_normalization_node)
# get Normalize
scale_shape, scale_data = self.GetParamsShapeAndData(Layers[i])
scale_shape = [1, scale_shape[0][0], 1, 1]
scale_input = self.AddInputsTVIFromParams(Layers[i], ["_scale"], [TensorProto.FLOAT])
mul_input_name = [lp_normalization_output_name[0], node_name + "_scale"]
mul_input_shape = [input_shape[0], scale_shape]
mul_node = op.create_mul_node(Layers[i], node_name + "_mul", mul_input_name, output_name,
mul_input_shape)
self.onnxNodeList.append(mul_node)
elif Layers[i].type == "Power":
# Power: Mul + Add + Pow
# create mul node
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
power, scale, shift = op.get_power_param(Layers[i])
scale_node_name = self.AddInputsTVIMannul(Layers[i], ["_scale"], [TensorProto.FLOAT], [np.shape(scale)], [scale])
mul_input_name = [input_name[0], scale_node_name[0]]
mul_node = op.create_mul_node(Layers[i], node_name + "_mul", mul_input_name, [output_name[0] + "_mul"],
[input_shape[0], np.shape(power)])
self.onnxNodeList.append(mul_node)
# create Add node
shift_param_name = self.AddInputsTVIMannul(Layers[i], ["_shift"], [TensorProto.FLOAT], [np.shape(scale)],
[shift])
add_input_name = [output_name[0] + "_mul", shift_param_name[0]]
add_node = op.create_add_node(Layers[i], node_name + "_add", add_input_name, [output_name[0] + "_add"], [input_shape[0], np.shape(shift)])
self.onnxNodeList.append(add_node)
# create Pow
power_param_name = self.AddInputsTVIMannul(Layers[i], ["_param_power"], [TensorProto.FLOAT], [np.shape(power)],[power])
power_input_name = [output_name[0] + "_add", power_param_name[0]]
power_node = op.create_power_node(Layers[i], node_name + "_power", power_input_name, output_name,
[input_shape[0], np.shape(power)])
self.onnxNodeList.append(power_node)
elif Layers[i].type == "TanH":
# 1.获取节点输入名、输入维度、输出名、节点名
input_name, input_shape = self.GetLastLayerOutNameAndShape(
Layers[i]) # 获取输入名列表和输入形状
output_name = self.GetCurrentLayerOutName(Layers[i]) # 获取输出名列表
node_name = Layers[i].name
# 2.构建tanh_node
tanh_node = op.createTanh(
Layers[i], node_name, input_name, output_name, input_shape)
# 3.添加节点到节点列表
self.onnxNodeList.append(tanh_node)
elif Layers[i].type == "Crop":
# Crop: Slice
# create Slice node
input_name, input_shape = self.GetLastLayerOutNameAndShape(Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
starts, ends, axes = op.get_crop_param(Layers[i],input_shape)
Crop_name=[]
Crop_name.append(input_name[0])
starts_param = self.AddInputsTVIMannul(Layers[i],
['_starts' + str(i)],
[TensorProto.INT64],
[np.shape(starts)],
[starts])
ends_param = self.AddInputsTVIMannul(Layers[i],
['_ends' + str(i)],
[TensorProto.INT64],
[np.shape(ends)], [ends])
axes_param = self.AddInputsTVIMannul(Layers[i],
['_axes' + str(i)],
[TensorProto.INT64],
[np.shape(axes)], [axes])
Crop_name.extend(starts_param)
Crop_name.extend(ends_param)
Crop_name.extend(axes_param)
crop_node = op.create_crop_node(Layers[i], node_name, Crop_name, output_name,
input_shape)
self.onnxNodeList.append(crop_node)
# MVN
elif Layers[i].type == "MVN":
# MVN: InstanceNormalization
# create InstanceNormalization
if Layers[i].mvn_param.normalize_variance == False or Layers[i].mvn_param.across_channels == True:
print("Failed type not support: " + Layers[i].type)
exit(-1)
input_name, input_shape = self.GetLastLayerOutNameAndShape(
Layers[i])
output_name = self.GetCurrentLayerOutName(Layers[i])
node_name = Layers[i].name
MVN_name = []
MVN_name.append(input_name[0])
scale, bias = op.get_InstanceNorm_param(Layers[i],input_shape)
scale_param = self.AddInputsTVIMannul(Layers[i],
['_scale' + str(i)],
[TensorProto.FLOAT],
[np.shape(scale)],
[scale])
bias_param = self.AddInputsTVIMannul(Layers[i],
['_bias' + str(i)],
[TensorProto.FLOAT],
[np.shape(bias)], [bias])
MVN_name.extend(scale_param)
MVN_name.extend(bias_param)
MVN_node = op.create_InstanceNorm_op(Layers[i], node_name,
MVN_name, output_name,
input_shape)
self.onnxNodeList.append(MVN_node)
else:
print("Failed type not support: " + Layers[i].type)
exit(-1)
# 判断当前节点是否是输出节点
def JudgeOutput(self, current_node, nodelist):
for output_name in current_node.outputs_name:
for node in nodelist:
if output_name in node.inputs_name:
return False
return True
# 添加模型输出信息和中间节点信息
def AddOutputsTVIAndValueInfo(self):
for i in range(len(self.onnxNodeList)):
if self.JudgeOutput(self.onnxNodeList[i], self.onnxNodeList): # 构建输出节点信息
lastnode = self.onnxNodeList[i]
for j in range(len(lastnode.outputs_shape)):
output_tvi = helper.make_tensor_value_info(lastnode.outputs_name[j], TensorProto.FLOAT,
lastnode.outputs_shape[j])
self.onnxmodel.addOutputsTVI(output_tvi)
else: # 构建中间节点信息
innernode = self.onnxNodeList[i]
for k in range(len(innernode.outputs_shape)):
hid_out_tvi = helper.make_tensor_value_info(innernode.outputs_name[k], TensorProto.FLOAT,
innernode.outputs_shape[k])
self.onnxmodel.addValueInfoTVI(hid_out_tvi)
# 创建模型
def createOnnxModel(self):
node_def = [Node.node for Node in self.onnxNodeList]
graph_def = helper.make_graph(
node_def,
self.onnxmodel.name,
self.onnxmodel.in_tvi,
self.onnxmodel.out_tvi,
self.onnxmodel.init_t,
value_info=self.onnxmodel.hidden_out_tvi
)
model_def = helper.make_model(graph_def, producer_name='Tencent YouTu')
print("2.onnx模型转换完成")
return model_def
from google.protobuf import text_format
from proto import caffe_upsample_pb2
import onnx
from onnx import utils
def LoadCaffeModel(net_path, model_path):
# read prototxt
net = caffe_upsample_pb2.NetParameter()
text_format.Merge(open(net_path).read(), net)
# read caffemodel
model = caffe_upsample_pb2.NetParameter()
f = open(model_path, 'rb')
model.ParseFromString(f.read())
f.close()
return net, model
def LoadOnnxModel(onnx_path):
onnxmodel = onnx.load(onnx_path)
return onnxmodel
def SaveOnnxModel(onnx_model, onnx_save_path, need_polish=True):
try:
if need_polish:
polished_model = onnx.utils.polish_model(onnx_model)
onnx.save_model(polished_model, onnx_save_path)
else:
onnx.save_model(onnx_model, onnx_save_path)
print("模型保存成功,已保存至:" + onnx_save_path)
except Exception as e:
print("模型存在问题,未保存成功:", e)
from proto import caffe_upsample_pb2
from onnx import TensorProto
Layer_CONCAT = caffe_upsample_pb2.V1LayerParameter.CONCAT # 3
Layer_CONVOLUTION = caffe_upsample_pb2.V1LayerParameter.CONVOLUTION # 4
Layer_DROPOUT = caffe_upsample_pb2.V1LayerParameter.DROPOUT # 6
Layer_INNER_PRODUCT = caffe_upsample_pb2.V1LayerParameter.INNER_PRODUCT # 14
Layer_LRN = caffe_upsample_pb2.V1LayerParameter.LRN # 15
Layer_POOLING = caffe_upsample_pb2.V1LayerParameter.POOLING # 17
Layer_RELU = caffe_upsample_pb2.V1LayerParameter.RELU # 18
Layer_SOFTMAX = caffe_upsample_pb2.V1LayerParameter.SOFTMAX # 20
Layer_ELTWISE = caffe_upsample_pb2.V1LayerParameter.ELTWISE # 25
Layer_UPSAMPLE = caffe_upsample_pb2.V1LayerParameter.UPSAMPLE # 40
op_pname = {"Conv": ["_W", "_b"],
"BatchNorm": ["_mean", "_var"],
"Scale": ["_scale", "_b"],
"Reshape": ["_shape"],
"DouReshape": ["_Doureshape"],
"InnerProduct": ["_W", "_B"],
"Upsample": ["_roi_" ,"_Scale"],
"PRelu": ["_slope"],
"Transpose": ["_trans"],
"ConvTranspose": ["_W", "_b"],
"Slice": ['_starts', '_ends', '_axes', '_steps']
}
op_ptype = {"Conv": [TensorProto.FLOAT, TensorProto.FLOAT],
"BatchNorm": [TensorProto.FLOAT, TensorProto.FLOAT],
"Scale": [TensorProto.FLOAT, TensorProto.FLOAT],
"Reshape": [TensorProto.INT64],
"InnerProduct": [TensorProto.FLOAT, TensorProto.FLOAT],
"Upsample": [TensorProto.FLOAT, TensorProto.FLOAT],
"PRelu": [TensorProto.FLOAT],
"Transpose": [TensorProto.INT64],
"ConvTranspose": [TensorProto.FLOAT, TensorProto.FLOAT],
"DouReshape": [TensorProto.INT64],
"Slice": [TensorProto.INT64, TensorProto.INT64, TensorProto.INT64, TensorProto.INT64]
}
# Tencent is pleased to support the open source community by making TNN available.
#
# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.
def is_ssd_model(proto_path):
proto_file = open(proto_path, 'r')
lines = proto_file.read()
proto_file.close()
if "PriorBox" in lines:
return True
elif "DetectionOutput" in lines:
return True
else:
return False
# Paddle 模型转换为 ONNX 模型
## 1. 环境搭建
- 安装PaddlePaddle
```shell script
pip install paddlepaddle
```
- 安装onnxruntime >= 1.10.0
```shell script
pip install onnxruntime
```
- 安装onnx
```
pip install onnx
```
- 安装paddle2onnx
```shell script
pip install paddle2onnx
```
## 2. caffe2onnx 工具使用
### 2.1获取PaddlePaddle部署模型
将paddle模型转换成onnx之前,请注意,所必需的模型文件:
- `model_name.pdmodel`: 表示模型结构
- `model_name.pdiparams`: 表示模型参数
[**注意**] 这里需要注意,两个文件其中参数文件后辍为 `.pdiparams`,如你的参数文件后辍是 `.pdparams` ,那说明你的参数是训练过程中保存的,当前还不是部署模型格式。
示例模型[下载地址](https://bj.bcebos.com/paddle2onnx/model_zoo/resnet50.tar.gz),下载模型后解压文件
```shell
tar -zxvf resnet50.tar.gz
```
## 命令行转换
```bash
paddle2onnx --model_dir resnet50 --model_filename inference.pdmodel --params_filename inference.pdiparams --save_file model.onnx
```
可调整的转换参数如下表,具体可参考[官网](https://github.com/PaddlePaddle/Paddle2ONNX):
| 参数 | 参数说明 |
| -------------------------- | ------------------------------------------------------------ |
| --model_dir | 配置包含 Paddle 模型的目录路径 |
| --model_filename | **[可选]** 配置位于 `--model_dir` 下存储网络结构的文件名 |
| --params_filename | **[可选]** 配置位于 `--model_dir` 下存储模型参数的文件名称 |
| --save_file | 指定转换后的模型保存目录路径 |
| --opset_version | **[可选]** 配置转换为 ONNX 的 OpSet 版本,目前支持 7~16 等多个版本,默认为 9 |
| --enable_onnx_checker | **[可选]** 配置是否检查导出为 ONNX 模型的正确性, 建议打开此开关, 默认为 False |
| --enable_auto_update_opset | **[可选]** 是否开启 opset version 自动升级功能,当低版本 opset 无法转换时,自动选择更高版本的 opset进行转换, 默认为 True |
| --deploy_backend | **[可选]** 量化模型部署的推理引擎,支持 onnxruntime、tensorrt 或 others,当选择 others 时,所有的量化信息存储于 max_range.txt 文件中,默认为 onnxruntime |
| --save_calibration_file | **[可选]** TensorRT 8.X版本部署量化模型需要读取的 cache 文件的保存路径,默认为 calibration.cache |
| --version | **[可选]** 查看 paddle2onnx 版本 |
| --external_filename | **[可选]** 当导出的 ONNX 模型大于 2G 时,需要设置 external data 的存储路径,推荐设置为:external_data |
| --export_fp16_model | **[可选]** 是否将导出的 ONNX 的模型转换为 FP16 格式,并用 ONNXRuntime-GPU 加速推理,默认为 False |
| --custom_ops | **[可选]** 将 Paddle OP 导出为 ONNX 的 Custom OP,例如:--custom_ops '{"paddle_op":"onnx_op"},默认为 {} |
## 修改模型shape
如果想要将输入name为image的模型model.onnx修改为shape[256, 3, 224, 224],可以用一下命令:
```bash
python -m paddle2onnx.optimize --input_model model.onnx --output_model new_model.onnx --input_shape_dict="{'inputs': [256, 3, 224, 224]}"
```
如果不知道模型的输入name,可以使用一下python脚本获知:
```shell
python get_model_input_output_info.py
```
其中get_model_input_output_info.py代码如下,替换model.onnx的路径即可查询自己的模型的输入输出信息。
```python
import onnx
def get_input_details(onnx_file):
# 加载ONNX模型
model = onnx.load(onnx_file)
# 获取输入信息
input_details = []
for input in model.graph.input:
input_name = input.name
input_shape = [dim.dim_value for dim in input.type.tensor_type.shape.dim]
input_details.append((input_name, input_shape))
return input_details
def get_output_details(onnx_file):
# 加载ONNX模型
model = onnx.load(onnx_file)
# 获取输出信息
output_details = []
for output in model.graph.output:
output_name = output.name
output_shape = [dim.dim_value for dim in output.type.tensor_type.shape.dim]
output_details.append((output_name, output_shape))
return output_details
# 调用函数并打印输入和输出信息
onnx_file = "model.onnx" # 将 "model.onnx" 替换为你的ONNX模型文件路径
input_details = get_input_details(onnx_file)
output_details = get_output_details(onnx_file)
print("输入信息:")
for name, shape in input_details:
print("Input Name:", name)
print("Input Shape:", shape)
print("输出信息:")
for name, shape in output_details:
print("Output Name:", name)
print("Output Shape:", shape)
```
\ No newline at end of file
import onnx
def get_input_details(onnx_file):
# 加载ONNX模型
model = onnx.load(onnx_file)
# 获取输入信息
input_details = []
for input in model.graph.input:
input_name = input.name
input_shape = [dim.dim_value for dim in input.type.tensor_type.shape.dim]
input_details.append((input_name, input_shape))
return input_details
def get_output_details(onnx_file):
# 加载ONNX模型
model = onnx.load(onnx_file)
# 获取输出信息
output_details = []
for output in model.graph.output:
output_name = output.name
output_shape = [dim.dim_value for dim in output.type.tensor_type.shape.dim]
output_details.append((output_name, output_shape))
return output_details
# 调用函数并打印输入和输出信息
onnx_file = "model.onnx" # 将 "your_model.onnx" 替换为你的ONNX模型文件路径
input_details = get_input_details(onnx_file)
output_details = get_output_details(onnx_file)
print("输入信息:")
for name, shape in input_details:
print("Input Name:", name)
print("Input Shape:", shape)
print("输出信息:")
for name, shape in output_details:
print("Output Name:", name)
print("Output Shape:", shape)
\ No newline at end of file
# tf转onnx
## 环境准备
- tensorflow安装
```bash
pip install tensorflow
```
- tf2onnx 安装(version>= 1.5.5)
```bash
pip install tf2onnx
```
## 模型格式确认
请先确认手里的模型文件的格式,一般情况下:
1. **SavedModel 文件结构**
- `saved_model.pb``saved_model.pbtxt`:这是SavedModel的核心文件,包含了模型的图(graph)和元数据(metadata)。
- `variables/`:这个文件夹包含两个文件,`variables.data-?????-of-?????``variables.index`,存储了模型的变量。
- `assets/`(可选):这个文件夹存储了任何附加的资源文件。
如果你的`.pb`文件位于一个包含上述结构的目录中,那么它很可能是一个SavedModel。
2. **Checkpoint 文件结构**
- Checkpoint 通常包含三个文件:一个`.index`文件,一个或多个`.data-?????-of-?????`文件,以及一个`checkpoint`文件,这个文件是保存模型变量的。
如果你的`.pb`文件位于一个包含上述结构的目录中,那么它很可能是一个Checkpoint。
3. **GraphDef 文件结构**
- 如果只有一个`.pb`文件,且没有与其相关联的其他文件或目录结构,那么它很可能是GraphDef。
可以使用一下代码对模型文件格式进行检查:
```python
import tensorflow as tf
from google.protobuf import text_format
from tensorflow.core.framework import graph_pb2
def is_saved_model(model_dir):
try:
model = tf.saved_model.load(model_dir)
return True
except Exception:
return False
def is_graph_def(pb_file):
try:
with tf.io.gfile.GFile(pb_file, "rb") as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
return True
except Exception:
return False
def is_checkpoint(model_dir):
try:
checkpoint = tf.train.Checkpoint()
checkpoint.restore(model_dir).expect_partial()
return True
except Exception:
return False
model_path = "/path/to/model"
if is_saved_model(model_path):
print(f"{model_path} contains a SavedModel.")
elif is_graph_def(model_path):
print(f"{model_path} contains a GraphDef.")
elif is_checkpoint(model_path):
print(f"{model_path} contains a Checkpoint.")
else:
print(f"{model_path} format is unknown.")
```
## 模型输入输出的name和shape确认
使用下面代码对GraphDef格式的模型进行确认
```python
import tensorflow as tf
from tensorflow.python.framework import tensor_util
def load_graph_def(pb_file):
with tf.io.gfile.GFile(pb_file, "rb") as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
return graph_def
def get_graph_inputs_outputs(graph_def):
inputs = []
outputs = []
for node in graph_def.node:
if node.op == 'Placeholder':
shape = None
for attr_value in node.attr.values():
if attr_value.HasField('shape'):
shape = [dim.size for dim in attr_value.shape.dim]
inputs.append({'name': node.name, 'shape': shape})
# Assuming outputs are nodes with no outputs themselves, usually not a strict rule
elif not any(node.name in input for input in [n.input for n in graph_def.node]):
shape = None
try:
tensor_shape = tensor_util.MakeNdarray(node.attr["shape"].shape)
shape = tensor_shape.shape
except:
pass
outputs.append({'name': node.name, 'shape': shape})
return inputs, outputs
def print_graph_info(inputs, outputs):
print("Inputs:")
for input_info in inputs:
print(f"Name: {input_info['name']}, Shape: {input_info['shape']}")
print("\nOutputs:")
for output_info in outputs:
print(f"Name: {output_info['name']}, Shape: {output_info['shape']}")
# Path to your .pb file
pb_file_path = "resnet50v15_tf.pb"
# Load GraphDef
graph_def = load_graph_def(pb_file_path)
# Get inputs and outputs
inputs, outputs = get_graph_inputs_outputs(graph_def)
# Print inputs and outputs
print_graph_info(inputs, outputs)
```
## 模型转换
使用tf2onnx工具进行模型转换,详细工具说明可以查看tf2onnx工具官网,tf2onnx项目地址:https://github.com/onnx/tensorflow-onnx 建议大家阅读 tf2onnx 的 README.md 文件,里面有详细的对该工具各个参数的说明。
```
options:
-h, --help show this help message and exit
--input INPUT input from graphdef
--graphdef GRAPHDEF input from graphdef
--saved-model SAVED_MODEL
input from saved model
--tag TAG tag to use for saved_model
--signature_def SIGNATURE_DEF
signature_def from saved_model to use
--concrete_function CONCRETE_FUNCTION
For TF2.x saved_model, index of func signature in __call__ (--signature_def is ignored)
--checkpoint CHECKPOINT
input from checkpoint
--keras KERAS input from keras model
--tflite TFLITE input from tflite model
--tfjs TFJS input from tfjs model
--large_model use the large model format (for models > 2GB)
--output OUTPUT output model file
--inputs INPUTS model input_names (optional for saved_model, keras, and tflite)
--outputs OUTPUTS model output_names (optional for saved_model, keras, and tflite)
--ignore_default IGNORE_DEFAULT
comma-separated list of names of PlaceholderWithDefault ops to change into Placeholder ops
--use_default USE_DEFAULT
comma-separated list of names of PlaceholderWithDefault ops to change into Identity ops using
their default value
--rename-inputs RENAME_INPUTS
input names to use in final model (optional)
--rename-outputs RENAME_OUTPUTS
output names to use in final model (optional)
--use-graph-names (saved model only) skip renaming io using signature names
--opset OPSET opset version to use for onnx domain
--dequantize remove quantization from model. Only supported for tflite currently.
--custom-ops CUSTOM_OPS
comma-separated map of custom ops to domains in format OpName:domain. Domain
'ai.onnx.converters.tensorflow' is used by default.
--extra_opset EXTRA_OPSET
extra opset with format like domain:version, e.g. com.microsoft:1
--load_op_libraries LOAD_OP_LIBRARIES
comma-separated list of tf op library paths to register before loading model
--target {rs4,rs5,rs6,caffe2,tensorrt,nhwc}
target platform
--continue_on_error continue_on_error
--verbose, -v verbose output, option is additive
--debug debug mode
--output_frozen_graph OUTPUT_FROZEN_GRAPH
output frozen tf graph to file
--inputs-as-nchw INPUTS_AS_NCHW
transpose inputs as from nhwc to nchw
--outputs-as-nchw OUTPUTS_AS_NCHW
transpose outputs as from nhwc to nchw
Usage Examples:
python -m tf2onnx.convert --saved-model saved_model_dir --output model.onnx
python -m tf2onnx.convert --input frozen_graph.pb --inputs X:0 --outputs output:0 --output model.onnx
python -m tf2onnx.convert --checkpoint checkpoint.meta --inputs X:0 --outputs output:0 --output model.onn
```
下面是将GraphDef格式的模型转换onnx的示例, resnet50v15_tf.pb模型[下载地址](https://obs-9be7.obs.cn-east-2.myhuaweicloud.com/003_Atc_Models/modelzoo/Official/cv/Resnet50v1.5_for_ACL/resnet50v15_tf.pb)
```bash
python -m tf2onnx.convert --graphdef resnet50v15_tf.pb --output model_nchw.onnx --inputs input_tensor:0 --outputs global_step:0,ArgMax:0,softmax_tensor:0 --inputs-as-nchw input_tensor:0
```
# tflite转onnx
同样使用tf2onnx工具,例如将ResNet50.tflite模型转为onnx模型,模型[下载地址](https://hf-mirror.com/qualcomm/ResNet50/resolve/main/ResNet50.tflite?download=true)
```bash
python -m tf2onnx.convert --opset 16 --tflite ResNet50.tflite --output model.onnx
```
\ No newline at end of file
import tensorflow as tf
from google.protobuf import text_format
from tensorflow.core.framework import graph_pb2
def is_saved_model(model_dir):
try:
model = tf.saved_model.load(model_dir)
return True
except Exception:
return False
def is_graph_def(pb_file):
try:
with tf.io.gfile.GFile(pb_file, "rb") as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
return True
except Exception:
return False
def is_checkpoint(model_dir):
try:
checkpoint = tf.train.Checkpoint()
checkpoint.restore(model_dir).expect_partial()
return True
except Exception:
return False
model_path = "predict_net.pb"
if is_saved_model(model_path):
print(f"{model_path} contains a SavedModel.")
elif is_graph_def(model_path):
print(f"{model_path} contains a GraphDef.")
elif is_checkpoint(model_path):
print(f"{model_path} contains a Checkpoint.")
else:
print(f"{model_path} format is unknown.")
# import tensorflow as tf
# from tensorflow.python.framework import graph_util
# from tensorflow.python.framework import graph_io
# def load_graph_def(pb_file):
# with tf.io.gfile.GFile(pb_file, "rb") as f:
# graph_def = tf.compat.v1.GraphDef()
# graph_def.ParseFromString(f.read())
# return graph_def
# def get_input_output_names(graph_def):
# input_names = []
# output_names = []
# for node in graph_def.node:
# if node.op == 'Placeholder':
# input_names.append(node.name)
# # Identify output nodes as those not used as inputs to other nodes
# is_output = True
# for n in graph_def.node:
# if node.name in n.input:
# is_output = False
# break
# if is_output:
# output_names.append(node.name)
# return input_names, output_names
# # 指定GraphDef pb文件路径
# pb_file_path = "resnet50v15_tf.pb"
# # 加载GraphDef
# graph_def = load_graph_def(pb_file_path)
# # 获取输入和输出的节点名称
# input_names, output_names = get_input_output_names(graph_def)
# print("Input names:", input_names)
# print("Output names:", output_names)
import tensorflow as tf
from tensorflow.python.framework import tensor_util
def load_graph_def(pb_file):
with tf.io.gfile.GFile(pb_file, "rb") as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
return graph_def
def get_graph_inputs_outputs(graph_def):
inputs = []
outputs = []
for node in graph_def.node:
if node.op == 'Placeholder':
shape = None
for attr_value in node.attr.values():
if attr_value.HasField('shape'):
shape = [dim.size for dim in attr_value.shape.dim]
inputs.append({'name': node.name, 'shape': shape})
# Assuming outputs are nodes with no outputs themselves, usually not a strict rule
elif not any(node.name in input for input in [n.input for n in graph_def.node]):
shape = None
try:
tensor_shape = tensor_util.MakeNdarray(node.attr["shape"].shape)
shape = tensor_shape.shape
except:
pass
outputs.append({'name': node.name, 'shape': shape})
return inputs, outputs
def print_graph_info(inputs, outputs):
print("Inputs:")
for input_info in inputs:
print(f"Name: {input_info['name']}, Shape: {input_info['shape']}")
print("\nOutputs:")
for output_info in outputs:
print(f"Name: {output_info['name']}, Shape: {output_info['shape']}")
# Path to your .pb file
pb_file_path = "resnet50v15_tf.pb"
# Load GraphDef
graph_def = load_graph_def(pb_file_path)
# Get inputs and outputs
inputs, outputs = get_graph_inputs_outputs(graph_def)
# Print inputs and outputs
print_graph_info(inputs, outputs)
# Pytorch模型转换onnx模型
## 环境准备
确保环境中有pytorch,可以用pip安装pytorch,具体可以参考PyTorch官网https://pytorch.org/get-started/locally/
```bash
pip install torch
```
## 模型确认
首先,请辨别手里的Pytorch模型是权重数据文件.pth还是包含模型结构以及权重数据的.pt文件,通常惯例是:
- **`.pth` 文件** 通常用于保存模型的权重(`state_dict`)。
- **`.pt` 文件** 通常用于保存整个模型(包括模型结构和权重)。
然而,这只是惯例,实际使用中两者可以互换,也取决于保存文件的人如何命名的,甚至可以是其他的后缀名。因此,为了准确加载模型,你需要知道该文件具体保存了什么。下面python脚本可以判断模型文件是否是完整模型文件。
```python
import torch
# 尝试加载 .pt 文件
try:
model_data = torch.load('resnet50.pt')
print(type(model_data)) # 打印数据类型
if isinstance(model_data, dict):
print(model_data.keys()) # 如果是字典,打印键
print("The .pt file is weights file)
else:
print("The .pt file contains the complete model.")
except Exception as e:
print(f"Error loading model: {e}")
```
## 导出模型
一般情况下,使用保存的权重文件,这里示例模型[下载地址](https://download.pytorch.org/models/resnet50-0676ba61.pth),转换模型用torch.onnx.export函数,具体操作如下。
其中,如果想导出动态模型,可以通过设置dynamic_axes,这里设置dynamic_axes={'input' : {0 : 'batch_size'}, 'output' : {0 : 'batch_size'}},即为将输入和输出tensor的名称为batch_size的索引为0的维度设置为动态模式,导出的模型输入shape将会变为[batch_size, 3, 224, 224]。同理,如果设置为dynamic_axes={'input' : {0 : 'batch_size', 2: 'height', 3: 'width'}, 'output' : {0 : 'batch_size', 2: 'height', 3: 'width'}},导出的模型输入输出shape将为变为[batch_size, 3, height, width]。
```python
import torch
import torchvision.models as models
# 定义模型结构(以 ResNet-50 为例)
model = models.resnet50() # 或者使用自己定义的模型实例
model.load_state_dict(torch.load('resnet50-0676ba61.pth'))
model.eval() # 设置为评估模式
# 示例输入
input_tensor = torch.randn(1, 3, 224, 224) # 修改为你的输入张量形状
# Export the model
torch.onnx.export(model, # 需要转换的pytorch模型变量
input_tensor, # 模型示例输入 (多输入情况为多变量tuple)
"resnet50.onnx", # 导出模型文件名
export_params=True, # 导出权重参数
opset_version=10, # 指定ONNX的opset版本(可选)
do_constant_folding=True, # 是否常量折叠(可选)
input_names = ['input'], # 模型输入的names
output_names = ['output'], # 模型输出的names
dynamic_axes={'input' : {0 : 'batch_size'}, # 动态维度指定
'output' : {0 : 'batch_size'}})
```
运行脚本即可得到onnx模型,如果想导出自己的onnx模型,自行做相应修改。
```bash
python export_onnx.py
```
import torch
import torchvision.models as models
# 定义模型结构(以 ResNet-50 为例)
model = models.resnet50() # 或者使用自己定义的模型实例
model.load_state_dict(torch.load('resnet50-0676ba61.pth'))
model.eval() # 设置为评估模式
# 示例输入
input_tensor = torch.randn(1, 3, 224, 224) # 修改为你的输入张量形状
# Export the model
torch.onnx.export(model, # 需要转换的pytorch模型变量
input_tensor, # 模型示例输入 (多输入情况为多变量tuple)
"resnet50.onnx", # 导出模型文件名
export_params=True, # 导出权重参数
opset_version=10, # 指定ONNX的opset版本(可选)
do_constant_folding=True, # 是否常量折叠(可选)
input_names = ['input'], # 模型输入的names
output_names = ['output'], # 模型输出的names
dynamic_axes={'input' : {0 : 'batch_size'}, # 动态维度指定
'output' : {0 : 'batch_size'}})
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment