diff --git a/README.md b/README.md index b81f2ce511d7ffe254eba24aeb07e6dd93f12e73..183fbd7a22e7ae0610c6fb7790bcdc7a9f6fb391 100644 --- a/README.md +++ b/README.md @@ -9,13 +9,13 @@ X2Paddle在多个主流的CV模型上,测试过TensorFlow/Caffe/ONNX模型的 ## 环境依赖 -python >= 3.5 +python == 2.7 | python >= 3.5 paddlepaddle >= 1.5.0 **按需安装以下依赖** tensorflow : tensorflow == 1.14.0 caffe : 无 -onnx : onnx == 1.5.0 pytorch == 1.1.0 +onnx : onnx == 1.5.0 onnxruntime == 0.4.0 ## 安装 ### 安装方式一(推荐) diff --git a/setup.py b/setup.py index 3470161b4d7620745cf4d7e206599447d7a399a3..0cf68ba9e2fa172fc835187b732e6df332709a04 100644 --- a/setup.py +++ b/setup.py @@ -23,4 +23,9 @@ setuptools.setup( "Operating System :: OS Independent", ], license='Apache 2.0', - entry_points={'console_scripts': ['x2paddle=x2paddle.convert:main']}) + entry_points={ + 'console_scripts': [ + 'x2paddle=x2paddle.convert:main', + 'onnx_infer=x2paddle.onnx_infer:main' + ] + }) diff --git a/x2paddle/__init__.py b/x2paddle/__init__.py index 98a433b31058a4c8166312666d91533a8e53e50a..3d187266f146fb216a04537da37dc5df0811acc4 100644 --- a/x2paddle/__init__.py +++ b/x2paddle/__init__.py @@ -1 +1 @@ -__version__ = "0.4.5" +__version__ = "0.5.0" diff --git a/x2paddle/convert.py b/x2paddle/convert.py index 000453f60c82ae67e6aaf01f2044dc522a518838..b4566d128666a92296b4c243b7f551e80a658f36 100644 --- a/x2paddle/convert.py +++ b/x2paddle/convert.py @@ -154,7 +154,7 @@ def onnx2paddle(model_path, save_dir): model = ONNXDecoder(model_path) from x2paddle.op_mapper.onnx_op_mapper import ONNXOpMapper - mapper = ONNXOpMapper(model) + mapper = ONNXOpMapper(model, save_dir) from x2paddle.optimizer.onnx_optimizer import ONNXOptimizer optimizer = ONNXOptimizer(mapper) diff --git a/x2paddle/decoder/onnx_decoder.py b/x2paddle/decoder/onnx_decoder.py index 959b5b12aeaea4c6a84aca84b79bfe712423d9b5..8955f1e0453904c42ebc5f6eb449aab201a18cbf 100644 --- a/x2paddle/decoder/onnx_decoder.py +++ b/x2paddle/decoder/onnx_decoder.py @@ -17,7 +17,6 @@ from x2paddle.core.fluid_code import FluidCode from onnx.checker import ValidationError from onnx.checker import check_model from onnx.utils import polish_model -from onnx.version_converter import convert_version from onnx import helper from onnx.helper import get_attribute_value, make_attribute from onnx.shape_inference import infer_shapes @@ -26,9 +25,11 @@ from onnx.numpy_helper import to_array from onnx import AttributeProto, TensorProto, GraphProto from collections import OrderedDict as Dict import onnx +from onnx.helper import ValueInfoProto import numpy as np from copy import deepcopy import logging as _logging +import os default_op_domain = 'ai.onnx' _logger = _logging.getLogger(__name__) @@ -47,6 +48,7 @@ class ONNXGraphNode(GraphNode): self.weight_inputs = list() self.out_shapes = list() self.dtype = None + self.which_child = {} def get_attr_map(self): """ @@ -60,10 +62,9 @@ class ONNXGraphNode(GraphNode): @property def value(self): assert 'Constant' in self.layer_type, "Only Constant | ConstantOfShape node has value." - attr = self.layer.attribute['value'] if 'value' not in self.attr_map: return None - return self.attr_map[name] + return self.attr_map['value'] def get_attribute_value2(self, attr): """ @@ -105,29 +106,39 @@ class ONNXGraphDataNode(GraphNode): self.fluid_code = FluidCode() self.weight = None self.embeded_as = None + self.which_child = {} @property def out_shapes(self): - values = self.layer.type.tensor_type.shape.dim - out_shapes = list() - out_shapes.append([dim.dim_value for dim in values]) - return out_shapes + if isinstance(self.layer, ValueInfoProto): + values = self.layer.type.tensor_type.shape.dim + out_shapes = list() + out_shapes.append([dim.dim_value for dim in values]) + return out_shapes + else: + values = self.layer.dims + out_shapes = list() + out_shapes.append(values) + return out_shapes @property def dtype(self): - dtype = self.layer.type.tensor_type.elem_type - return TENSOR_TYPE_TO_NP_TYPE[dtype] + if isinstance(self.layer, ValueInfoProto): + dtype = self.layer.type.tensor_type.elem_type + return TENSOR_TYPE_TO_NP_TYPE[dtype] + else: + dtype = self.layer.data_type + return TENSOR_TYPE_TO_NP_TYPE[dtype] class ONNXGraph(Graph): - def __init__(self, graph, onnx_model): - super(ONNXGraph, self).__init__(graph) + def __init__(self, onnx_model): + super(ONNXGraph, self).__init__(onnx_model.graph) self.onnx_model = onnx_model self.initializer = {} self.place_holder_nodes = list() self.get_place_holder_nodes() - - self.value_infos = self.inferred_model_value_info(graph) + self.value_infos = self.inferred_model_value_info(self.model) self.results_of_inference = dict() def get_inner_nodes(self): @@ -165,22 +176,9 @@ class ONNXGraph(Graph): """ build topo_sort of ONNX model """ - data_node = self.place_holder_nodes[0] - value_info = self.value_infos[data_node] - input_shape = value_info['shape'] - self.get_results_of_inference(self.onnx_model, input_shape) for layer in self.model.node: node = ONNXGraphNode(layer) self.node_map[layer.name] = node - for opt in layer.output: - if opt in self.value_infos: - value_info = self.value_infos[opt] - node.dtype = value_info['dtype'] - node.out_shapes.append(value_info['shape']) - else: - _, dtype, shape = self.get_dynamic_shape(opt) - node.dtype = dtype - node.out_shapes.append(shape) for layer in self.model.input: if layer.name not in self.node_map: @@ -191,20 +189,40 @@ class ONNXGraph(Graph): is_global_input=is_place_holder) #set data node's weight - for name, weight in self.graph_weights(self.model): + for initializer in self.model.initializer: + name = initializer.name + weight = to_array(initializer) if name in self.node_map: if isinstance(self.node_map[name], ONNXGraphDataNode): self.node_map[name].weight = weight self.node_map[name].embeded_as = [] + else: + self.node_map[name] = ONNXGraphDataNode(initializer, + layer_name=name, + is_global_input=False) + self.node_map[name].weight = weight + self.node_map[name].embeded_as = [] #generate connection between nodes for topo for layer_name, node in self.node_map.items(): if isinstance(node, ONNXGraphNode): for idx, in_node in enumerate(node.layer.input): if in_node not in self.node_map: - raise Exception( - 'input[{}] of node[{}] does not exist in node_map'. - format(in_node, layer_name)) + flag = 0 + for nd in self.model.node: + for idx, opt in enumerate(nd.output): + if opt == in_node: + self.connect(nd.name, layer_name) + flag = 1 + node.which_child[nd.name] = idx + self.node_map[nd.name].index = 0 + break + if flag == 1: + break + if flag == 0: + raise Exception( + 'input[{}] of node[{}] does not exist in node_map' + .format(in_node, layer_name)) else: self.connect(in_node, layer_name) #generate topo @@ -212,13 +230,16 @@ class ONNXGraph(Graph): self.input_nodes = self.place_holder_nodes - def get_nodes(self, names, copy=False): - """ - get nodes by more than one name - """ - nodes = [] - for name in names: - nodes.add(self.get_node(name, copy=copy)) + def get_input_node(self, node, idx=0, copy=False): + if len(node.which_child) == 0: + ipt_node = super(ONNXGraph, self).get_node(node.inputs[idx], copy) + return ipt_node + + else: + ipt_node = super(ONNXGraph, self).get_node(node.inputs[idx], copy) + if ipt_node.layer_name in node.which_child: + ipt_node.index = node.which_child[ipt_node.layer_name] + return ipt_node def graph_weights(self, graph): """ @@ -270,50 +291,6 @@ class ONNXGraph(Graph): } return value_info - def get_results_of_inference(self, model, shape): - try: - import torch - version = torch.__version__ - if '1.1.0' not in version: - print("your model have dynamic graph, torch==1.1.0 is required") - return - except: - print( - "your model have dynamic graph, we use caff2 to inference graph, please use \"pip install torch==1.1.0\"." - ) - return - from x2paddle.decoder.onnx_backend import prepare - - np_images = np.random.rand(shape[0], shape[1], shape[2], - shape[3]).astype('float32') - - outputs = [] - for node in model.graph.node: - value_info = helper.make_tensor_value_info(node.name, - TensorProto.UNDEFINED, - []) - outputs.append(value_info) - - while len(outputs) > 0: - tmp_outputs = outputs[:254] - model.graph.ClearField('output') - model.graph.output.MergeFrom(tmp_outputs) - prepared_backend = prepare(model, - device='CPU', - no_check_UNSAFE=True) - res = prepared_backend.run(inputs=np_images) - for idx, info in enumerate(tmp_outputs): - self.results_of_inference[info.name] = res[idx] - outputs = outputs[254:] - return - - def get_dynamic_shape(self, layer): - """ - get dynamic shape from caffe2.backend - """ - output = self.results_of_inference[layer] - return output.tolist(), output.dtype, output.shape - class ONNXDecoder(object): def __init__(self, onnx_model): @@ -334,8 +311,8 @@ class ONNXDecoder(object): self.standardize_variable_name(model.graph) self.model = model - graph_def = model.graph - self.onnx_graph = ONNXGraph(graph_def, model) + graph = model.graph + self.onnx_graph = ONNXGraph(model) self.onnx_graph.build() def build_value_refs(self, nodes): @@ -476,7 +453,7 @@ class ONNXDecoder(object): if name == '': raise ValueError('name should not be empty') - for s in ' .*?\\/-:': # + for s in ' .*?\\/-:': name = name.replace(s, '_') return '_' + name @@ -499,46 +476,3 @@ class ONNXDecoder(object): node.input[i] = self.make_variable_name(node.input[i]) for i in range(len(node.output)): node.output[i] = self.make_variable_name(node.output[i]) - - def split_model(self, model, outputs=None): - """ - Takes a model and changes its outputs. - """ - if outputs is None: - raise RuntimeError("outputs is None") - if outputs == model.graph.output[0].name: - return model - nodes = model.graph.node - keep_nodes = [] - - # all the nodes we need to keep. - for node in nodes: - if outputs in node.output: - keep_nodes.append(node) - break - keep_nodes.append(node) - - infer_shapes = onnx.shape_inference.infer_shapes(model) - - var_out = [] - for value_info in infer_shapes.graph.value_info: - if value_info.name == outputs: - var_out.append(value_info) - break - - graph = helper.make_graph(keep_nodes, model.graph.name, - model.graph.input, var_out, - model.graph.initializer) - - onnx_model = helper.make_model(graph) - onnx_model.ir_version = model.ir_version - onnx_model.producer_name = model.producer_name - onnx_model.producer_version = model.producer_version - onnx_model.domain = model.domain - onnx_model.model_version = model.model_version - onnx_model.doc_string = model.doc_string - - if len(onnx_model.graph.input) != len(model.graph.input): - raise RuntimeError("Input mismatch {} != {}".format( - len(onnx_model.input), len(model.input))) - return onnx_model diff --git a/x2paddle/onnx_infer.py b/x2paddle/onnx_infer.py new file mode 100644 index 0000000000000000000000000000000000000000..8d6b0f222efa2cebff900377f1293b1aa7ed05ef --- /dev/null +++ b/x2paddle/onnx_infer.py @@ -0,0 +1,54 @@ +import os +import sys +import numpy as np +import onnx +import json +import argparse +from six import text_type as _text_type + + +def arg_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("--save_dir", + "-s", + type=_text_type, + default=None, + help="define save_dir") + return parser + + +def main(): + try: + import onnxruntime as rt + version = rt.__version__ + if version != '0.4.0': + print("onnxruntime==0.4.0 is required") + return + except: + print( + "onnxruntime is not installed, use \"pip install onnxruntime==0.4.0\"." + ) + return + parser = arg_parser() + args = parser.parse_args() + + save_dir = args.save_dir + model_dir = os.path.join(save_dir, 'onnx_model_infer.onnx') + data_dir = os.path.join(save_dir, 'input_data.npy') + + model = onnx.load(model_dir) + sess = rt.InferenceSession(model_dir) + + inputs = np.load(data_dir, allow_pickle=True) + data_dir + inputs_dict = {} + for i, ipt in enumerate(inputs): + inputs_dict[sess.get_inputs()[i].name] = ipt + + res = sess.run(None, input_feed=inputs_dict) + for idx, value_info in enumerate(model.graph.output): + np.save(os.path.join(save_dir, value_info.name), res[idx]) + + +if __name__ == "__main__": + main() diff --git a/x2paddle/op_mapper/caffe_op_mapper.py b/x2paddle/op_mapper/caffe_op_mapper.py index 1d451639e7063b72b888ecf913c9d2be9cc8bb3e..3e3f249f2e75ace83876d5154aa1f9a034f1b668 100644 --- a/x2paddle/op_mapper/caffe_op_mapper.py +++ b/x2paddle/op_mapper/caffe_op_mapper.py @@ -954,6 +954,13 @@ class CaffeOpMapper(OpMapper): inputs_node = [] for i in range(len(node.inputs)): input = self.graph.get_bottom_node(node, idx=i, copy=True) + if i == 1 and op == 'DetectionOutput': + input = self.graph.get_bottom_node(node, idx=i, copy=True) + print(input.layer_type) + while input is not None and input.layer_type != 'Softmax': + input = self.graph.get_bottom_node(input, idx=0, copy=True) + assert input is not None, 'This kind of DetectionOutput is not supported!' + input = self.graph.get_bottom_node(input, idx=0, copy=True) inputs_node.append(input) node.fluid_code.add_layer(func.__code__.co_name, inputs=inputs_node, diff --git a/x2paddle/op_mapper/onnx_custom_layer/InstanceNormalization.py b/x2paddle/op_mapper/onnx_custom_layer/InstanceNormalization.py index f93f7a723f17c927a26dc12a5f20a27f112546f2..0d767d70b373d1316d92a6aee6bc3f628c9d9530 100644 --- a/x2paddle/op_mapper/onnx_custom_layer/InstanceNormalization.py +++ b/x2paddle/op_mapper/onnx_custom_layer/InstanceNormalization.py @@ -22,8 +22,9 @@ def InstanceNormalization_shape(input_shape): def InstanceNormalization_layer(inputs, name=None): # TODO(lvmengsi@baidu.com): Check the accuracy when using fluid.layers.layer_norm. epsilon = 1e-5 - mean = fluid.layers.reduce_mean(inputs, dim=[2, 3], keep_dim=True) - var = fluid.layers.reduce_mean(fluid.layers.square(inputs - mean), + input_ = inputs[0] + mean = fluid.layers.reduce_mean(input_, dim=[2, 3], keep_dim=True) + var = fluid.layers.reduce_mean(fluid.layers.square(input_ - mean), dim=[2, 3], keep_dim=True) if name is not None: @@ -36,13 +37,13 @@ def InstanceNormalization_layer(inputs, name=None): initializer=fluid.initializer.Constant(0.0), trainable=True) scale = fluid.layers.create_parameter(attr=scale_param, - shape=inputs.shape[1:2], + shape=input_.shape[1:2], dtype="float32") offset = fluid.layers.create_parameter(attr=offset_param, - shape=inputs.shape[1:2], + shape=input_.shape[1:2], dtype="float32") - tmp = fluid.layers.elementwise_mul(x=(inputs - mean), y=scale, axis=1) + tmp = fluid.layers.elementwise_mul(x=(input_ - mean), y=scale, axis=1) tmp = tmp / fluid.layers.sqrt(var + epsilon) tmp = fluid.layers.elementwise_add(tmp, offset, axis=1) return tmp @@ -56,4 +57,5 @@ def InstanceNormalization_weights(name, data=None): register(kind='InstanceNormalization', shape=InstanceNormalization_shape, layer=InstanceNormalization_layer, + child_func=None, weights=InstanceNormalization_weights) diff --git a/x2paddle/op_mapper/onnx_custom_layer/__init__.py b/x2paddle/op_mapper/onnx_custom_layer/__init__.py index c21ce95ef487405e9cd93300b5e3831fe79c6dd4..3c074350a63e9e24cf1eba441da5a420868b5bb7 100644 --- a/x2paddle/op_mapper/onnx_custom_layer/__init__.py +++ b/x2paddle/op_mapper/onnx_custom_layer/__init__.py @@ -95,6 +95,17 @@ def make_custom_layer(node): return inspect.getsource(layer_func), layer_func +def make_custom_child_func(node): + """ get the code which implement the custom layer function + """ + layer_type = node.layer_type + child_func = custom_layers[layer_type]['child_func'] + if child_func is None: + return None, child_func + import inspect + return inspect.getsource(child_func), child_func + + def deal_weights(node, data=None): """ deal the weights of the custom layer """ diff --git a/x2paddle/op_mapper/onnx_custom_layer/register.py b/x2paddle/op_mapper/onnx_custom_layer/register.py index ca3643c0497bba1a77fad8c2b819bbba08b8ba05..fb26c369f49bcf4c46bf0feba98a4345b092c6bd 100644 --- a/x2paddle/op_mapper/onnx_custom_layer/register.py +++ b/x2paddle/op_mapper/onnx_custom_layer/register.py @@ -17,7 +17,7 @@ g_custom_layers = {} -def register(kind, shape, layer, weights): +def register(kind, shape, layer, child_func, weights): """ register a custom layer or a list of custom layers Args: @@ -48,6 +48,7 @@ def register(kind, shape, layer, weights): g_custom_layers[k] = { 'shape': shape, 'layer': layer, + 'child_func': child_func, 'weights': weights } diff --git a/x2paddle/op_mapper/onnx_directly_map.py b/x2paddle/op_mapper/onnx_directly_map.py index bf26dbd911c3fb824f7cd7308a9a743220e3b4a0..48963d1f4e2d9f833c9747ae32f8d9b5b7990dae 100644 --- a/x2paddle/op_mapper/onnx_directly_map.py +++ b/x2paddle/op_mapper/onnx_directly_map.py @@ -29,9 +29,6 @@ default_op_mapping = { 'Gather': ['gather', ['X'], ['Out'], dict(axis='')], 'Shape': ['shape', ['X'], ['Out']], - 'Mul': ['elementwise_mul', ['X', 'Y'], ['Out'], - dict(), - dict(axis=-1)], 'Clip': [ 'clip', ['X'], ['Out'], dict(), @@ -42,6 +39,7 @@ default_op_mapping = { dtype=_np.uint8).view(_np.float32)), ) ], + 'Ceil': ['ceil', ['X'], ['Out']], 'ReduceMean': [ 'reduce_mean', ['X'], ['Out'], dict(axes='dim', keepdims='keep_dim'), @@ -52,7 +50,11 @@ default_op_mapping = { dict(axes='dim', keepdims='keep_dim'), dict(keep_dim=1) ], - + 'ReduceMin': [ + 'reduce_min', ['X'], ['Out'], + dict(axes='dim', keepdims='keep_dim'), + dict(keep_dim=1) + ], #active function 'Relu': ['relu', ['X'], ['Out']], 'LeakyRelu': ['leaky_relu', ['X'], ['Out'], @@ -66,9 +68,6 @@ default_op_mapping = { ], 'Tanh': ['tanh', ['X'], ['Out']], 'Sigmoid': ['sigmoid', ['X'], ['Out']], - 'Pow': ['elementwise_pow', ['X', 'Y'], ['Out'], - dict(), - dict(axis=-1)], # TODO: pow for scalar exponent 'HardSigmoid': [ 'hard_sigmoid', ['X'], ['Out'], dict(alpha='slope', beta='offset'), @@ -78,8 +77,8 @@ default_op_mapping = { 'Softplus': ['softplus', ['X'], ['Out']], 'Exp': ['exp', ['X'], ['Out']], 'Softmax': ['softmax', ['X'], ['Out'], - dict(axis=''), - dict(axis=1)], + dict(), dict(axis=1)], + 'Sqrt': ['sqrt', ['X'], ['Out']], } activefunc_op_mapping = { diff --git a/x2paddle/op_mapper/onnx_op_mapper.py b/x2paddle/op_mapper/onnx_op_mapper.py index 4dcc922c15d21612d409dd2e568ca83335c7e194..f8382ab6811e282d882b31a86b738eb6776b4ff7 100644 --- a/x2paddle/op_mapper/onnx_op_mapper.py +++ b/x2paddle/op_mapper/onnx_op_mapper.py @@ -23,16 +23,21 @@ from x2paddle.op_mapper.onnx_directly_map import default_ioa_constraint from x2paddle.op_mapper.onnx_custom_layer import * from x2paddle.core.util import string import numpy as np +import onnx import onnx.numpy_helper as numpy_helper +from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE import logging as _logging from collections import OrderedDict as _dict +import math +import os +import shutil _logger = _logging.getLogger(__name__) def _const_weight_or_none(node): if 'Constant' in node.layer_name: - return val.value + return node.value if isinstance(node, ONNXGraphDataNode): return node.weight return None @@ -47,7 +52,7 @@ def get_same_padding(in_size, kernel_size, stride): class ONNXOpMapper(OpMapper): - def __init__(self, decoder): + def __init__(self, decoder, save_dir): super(ONNXOpMapper, self).__init__() self.decoder = decoder self.graph = decoder.onnx_graph @@ -55,6 +60,9 @@ class ONNXOpMapper(OpMapper): self.weights = dict() self.omit_nodes = list() self.used_custom_layers = dict() + self.is_inference = False + self.tmp_data_dir = os.path.join(save_dir, 'tmp_data') + self.get_output_shapes() if not self.op_checker(): raise Exception("Model are not supported yet.") @@ -76,6 +84,8 @@ class ONNXOpMapper(OpMapper): elif op in custom_layers: self.deal_custom_layer(node) + self.remove_tmp_data() + def op_checker(self): unsupported_ops = set() for node_name in self.graph.topo_sort: @@ -94,12 +104,85 @@ class ONNXOpMapper(OpMapper): print(op) return False + def get_results_of_inference(self, model, value_infos, data_nodes): + inputs = [] + for data_node in data_nodes: + value_info = value_infos[data_node] + ipt = np.random.random(value_info['shape']).astype( + value_info['dtype']) + inputs.append(ipt) + + model = onnx.shape_inference.infer_shapes(model) + outputs = [] + for value_info in model.graph.value_info: + outputs.append(value_info) + + model.graph.ClearField('output') + model.graph.output.MergeFrom(outputs) + if not os.path.exists(self.tmp_data_dir): + os.makedirs(self.tmp_data_dir) + onnx.save(model, os.path.join(self.tmp_data_dir, + 'onnx_model_infer.onnx')) + np.save(os.path.join(self.tmp_data_dir, 'input_data.npy'), inputs) + os.system('onnx_infer --save_dir=' + self.tmp_data_dir) + return + + def get_dynamic_shape(self, layer): + """ + get dynamic shape from infer_result + """ + output = np.load(os.path.join(self.tmp_data_dir, layer + '.npy')) + return output.tolist(), output.dtype, output.shape + + def get_output_shapes(self): + """ + build topo_sort of ONNX model + """ + nodes = self.decoder.model.graph.node + node_map = self.decoder.onnx_graph.node_map + value_infos = self.decoder.onnx_graph.value_infos + onnx_model = self.decoder.model + for layer in nodes: + node = node_map[layer.name] + for opt in layer.output: + if opt in value_infos: + value_info = value_infos[opt] + if len(value_info['shape'] + ) == 0 or value_info['dtype'] is None: + if self.is_inference == False: + self.get_results_of_inference( + onnx_model, value_infos, + self.decoder.onnx_graph.place_holder_nodes) + self.is_inference = True + _, dtype, shape = self.get_dynamic_shape(opt) + node.out_shapes.append(shape) + node.dtype = dtype + else: + node.dtype = value_info['dtype'] + node.out_shapes.append(value_info['shape']) + else: + if self.is_inference == False: + self.get_results_of_inference( + onnx_model, value_infos, + self.decoder.onnx_graph.place_holder_nodes) + self.is_inference = True + _, dtype, shape = self.get_dynamic_shape(opt) + node.dtype = dtype + node.out_shapes.append(shape) + + def remove_tmp_data(self): + """ + remove temporarily generated file + """ + if os.path.exists(self.tmp_data_dir): + import shutil + shutil.rmtree(self.tmp_data_dir) + def directly_map(self, node, name='', *args, **kwargs): inputs = node.layer.input outputs = node.layer.output op_type = node.layer_type attrs = node.attr_map - info = default_op_mapping[op_type] info.extend(list(default_op_mapping_field_values.values())[len(info):]) ( @@ -127,34 +210,41 @@ class ONNXOpMapper(OpMapper): mapped_attrs.pop('_') fluid_attrs = default_attrs.copy() fluid_attrs.update(mapped_attrs) - val_inps = inputs if input_perm is None else list( + inputs = inputs if input_perm is None else list( map(lambda i: inputs[i], input_perm)) + val_inps = [] + for idx, ipt in enumerate(inputs): + val_inps.append(self.graph.get_input_node(node, idx=idx, copy=True)) + val_outs = outputs if output_perm is None else list( map(lambda i: outputs[i], output_perm)) attr = fluid_attrs - if fluid_op not in ['shape', 'gather']: + assert len(val_inps) == 1, 'directly_map error with multi inputs' + if fluid_op not in ['shape']: attr['name'] = string(node.layer_name) node.fluid_code.add_layer(fluid_op, - inputs=', '.join(val_inps), + inputs=val_inps[0], output=val_outs[0], param_attr=attr) def deal_custom_layer(self, node): op = node.layer_type - val_x = self.graph.get_node(node.layer.input[0], copy=True) custom_code, func = make_custom_layer(node) + child_func_code, child_func = make_custom_child_func(node) params = get_params(node.layer, node.layer_type) arg_names, kwargs = set_args(func, params) kwargs['name'] = string(node.layer_name) - inputs_node = [] - inputs_node.append(node.inputs[0]) node.fluid_code.add_layer(func.__code__.co_name, - inputs=inputs_node[0], + inputs=node.inputs, output=node, param_attr=kwargs, is_custom_layer=True) if op not in self.used_custom_layers: self.used_custom_layers[op] = custom_code + if op + '_child_func' not in self.used_custom_layers: + if child_func_code is not None: + self.used_custom_layers[op + + '_child_func'] = child_func_code def place_holder(self, node): self.input_shapes.append(node.out_shapes[0]) @@ -203,8 +293,8 @@ class ONNXOpMapper(OpMapper): return [0] * ndims, val_padded def _interpolate(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_scales = self.graph.get_node(node.layer.input[1], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_scales = self.graph.get_input_node(node, idx=1, copy=True) val_y = self.graph.get_node(node.layer.output[0], copy=True) out_shape_ = val_y.out_shapes[0] @@ -245,7 +335,7 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Pad(self, node, op_independent=True): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) pads = node.get_attr('pads') mode = node.get_attr('mode', 'constant') value = node.get_attr('value', 0.) @@ -292,7 +382,7 @@ class ONNXOpMapper(OpMapper): return node.layer_name + '_paded' def Unsqueeze(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) axes = node.get_attr('axes') attr = {'axes': axes, 'name': string(node.layer_name)} node.fluid_code.add_layer('unsqueeze', @@ -301,7 +391,7 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Shrink(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) bias = node.get_attr('bias') lambd = node.get_attr('lambd') assert bias == 0.0, 'not support bias!=0' @@ -358,8 +448,8 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Resize(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_scales = self.graph.get_node(node.layer.input[1], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_scales = self.graph.get_input_node(node, idx=1, copy=True) val_y = self.graph.get_node(node.layer.output[0], copy=True) out_shape_ = val_y.out_shapes[0] @@ -401,24 +491,76 @@ class ONNXOpMapper(OpMapper): def Upsample(self, node): self._interpolate(node) + def Gather(self, node): + val_x = self.graph.get_input_node(node, idx=0, copy=True) + indices = self.graph.get_input_node(node, idx=1, copy=True) + indices_shape = indices.out_shapes[0] + axis = node.get_attr('axis') + assert len( + indices_shape) <= 1, "Gather op don't support dim of indice >1 " + if axis == 0 and len(indices_shape) <= 1: + node.fluid_code.add_layer('gather', + inputs={ + 'input': val_x, + 'index': indices + }, + output=node, + param_attr=None) + elif axis > 0 and len(indices_shape) <= 1: + perm = list(range(len(val_x.out_shapes[0]))) + perm = [axis] + perm[:axis] + perm[axis + 1:] + attr_trans = {'perm': perm} + name_trans = val_x.layer_name + '_trans' + node.fluid_code.add_layer('transpose', + inputs=val_x, + output=name_trans, + param_attr=attr_trans) + node.fluid_code.add_layer('gather', + inputs={ + 'input': name_trans, + 'index': indices + }, + output=node, + param_attr=None) + node.fluid_code.add_layer('transpose', + inputs=node, + output=node, + param_attr=attr_trans) + def Slice(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_starts, val_ends, val_axes, val_steps = None, None, None, None + if len(node.inputs) > 1: + starts = self.graph.get_input_node(node, idx=1, copy=True) + ends = self.graph.get_input_node(node, idx=2, copy=True) + axes = self.graph.get_input_node(node, idx=3, copy=True) + steps = self.graph.get_input_node(node, idx=4, copy=True) + + self.omit_nodes.append(starts.layer_name) + self.omit_nodes.append(ends.layer_name) + self.omit_nodes.append(axes.layer_name) + self.omit_nodes.append(steps.layer_name) + + starts = _const_weight_or_none(starts).copy() + ends = _const_weight_or_none(ends).copy() + axes = _const_weight_or_none(axes) + steps = _const_weight_or_none(steps) + else: + starts = node.get_attr('starts') + ends = node.get_attr('ends') + axes = node.get_attr('axes') + val_y = self.graph.get_node(node.layer.output[0], copy=True) - axes = node.get_attr('axes') - starts = node.get_attr('starts') - ends = node.get_attr('ends') shape = val_x.out_shapes[0] if shape is not None: for idx, value in enumerate(starts): - if value > 2**63 - 1 // 2: - value = value - ONNX_INT_MAX - starts[idx] = shape[axes[idx]] + value + if value > shape[axes[idx]]: + starts[idx] = shape[axes[idx]] for idx, value in enumerate(ends): - if value > 2**63 - 1 // 2: - value = value - ONNX_INT_MAX - ends[idx] = shape[axes[idx]] + value + if value > shape[axes[idx]]: + ends[idx] = shape[axes[idx]] attr = {"axes": axes, "starts": starts, "ends": ends} node.fluid_code.add_layer('slice', inputs=val_x, @@ -426,7 +568,7 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def ConstantOfShape(self, node): - val_shape = self.graph.get_node(node.layer.input[0], copy=True) + val_shape = self.graph.get_input_node(node, idx=0, copy=True) val_y = self.graph.get_node(node.layer.output[0], copy=True) shape = _const_weight_or_none(val_shape) @@ -452,31 +594,36 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Split(self, node): - val_input = self.graph.get_node(node.layer.input[0], copy=True) - var_outs = [val for val in node.layer.input] + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_y = self.graph.get_node(node.layer.output[0], copy=True) fluid_op = 'split' - split = node.get_attr['split'] + split = node.get_attr('split') axis = node.get_attr('axis', 0) - attr = {'split': split, 'axis': axis, 'name': string(node.layer_name)} + attr = { + 'num_or_sections': split, + 'dim': axis, + 'name': string(node.layer_name) + } # generation node.fluid_code.add_layer('split', - inputs=val_input, - output=var_outs, + inputs=val_x, + output=val_y, param_attr=attr) def Reshape(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_shape = self.graph.get_node(node.layer.input[1], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_shape = self.graph.get_input_node(node, idx=1, copy=True) val_reshaped = self.graph.get_node(node.layer.output[0], copy=True) shape = None + if isinstance(val_shape, ONNXGraphDataNode): self.omit_nodes.append(val_shape.layer_name) # catch dynamic graph shape if isinstance(val_shape, ONNXGraphNode): - shape, _, _ = self.decoder.onnx_graph.get_dynamic_shape( - val_shape.layer_name) + shape, _, _ = self.get_dynamic_shape(val_shape.layer_name) + if shape is None: shape = val_reshaped.out_shapes[0] @@ -503,7 +650,7 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Cast(self, node): - val_input = self.graph.get_node(node.layer.input[0], copy=True) + val_input = self.graph.get_input_node(node, idx=0, copy=True) val_output = self.graph.get_node(node.layer.output[0], copy=True) dtype = node.get_attr('to') @@ -520,7 +667,7 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def AveragePool(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) auto_pad = node.get_attr('auto_pad', 'NOTSET') kernel_shape = node.get_attr("kernel_shape") @@ -532,10 +679,10 @@ class ONNXOpMapper(OpMapper): fluid_op = 'pool{}d'.format(poolnd) assert 2 <= poolnd <= 3, 'only pool2d and pool3d is supported' - input_shape = val_x.out_shapes[0] paddings, val_x = self._pad_if_asymmetric(node, pads, val_x) if auto_pad == "SAME_UPPER" or auto_pad == "SAME_LOWER": + input_shape = val_x.out_shapes[0] pad_h = get_same_padding(input_shape[2], kernel_shape[0], strides[0]) pad_w = get_same_padding(input_shape[3], kernel_shape[1], @@ -560,7 +707,7 @@ class ONNXOpMapper(OpMapper): def Concat(self, node): inputs = [] for i in range(len(node.layer.input)): - ipt = self.graph.get_node(node.layer.input[i], copy=True) + ipt = self.graph.get_input_node(node, idx=i, copy=True) if isinstance(ipt, str): inputs.append(ipt) else: @@ -568,12 +715,12 @@ class ONNXOpMapper(OpMapper): axis = node.get_attr('axis') attr = {'axis': axis} node.fluid_code.add_layer('concat', - inputs='[' + ', '.join(inputs) + ']', + inputs=inputs, output=node, param_attr=attr) def Flatten(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) axis = node.get_attr('axis', 1) attr = {"axis": str(axis), "name": string(node.layer_name)} node.fluid_code.add_layer('flatten', @@ -582,9 +729,9 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Gemm(self, node): - val_a = self.graph.get_node(node.layer.input[0], copy=True) - val_b = self.graph.get_node(node.layer.input[1], copy=True) - val_c = self.graph.get_node(node.layer.input[2], copy=True) + val_a = self.graph.get_input_node(node, idx=0, copy=True) + val_b = self.graph.get_input_node(node, idx=1, copy=True) + val_c = self.graph.get_input_node(node, idx=2, copy=True) alpha = node.get_attr('alpha', 1.) # optional beta = node.get_attr('beta', 1.) # optional @@ -627,8 +774,8 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Add(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_y = self.graph.get_node(node.layer.input[1], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_y = self.graph.get_input_node(node, idx=1, copy=True) inputs = { "x": val_x, "y": val_y, @@ -639,26 +786,53 @@ class ONNXOpMapper(OpMapper): output=node, param_attr=attr) + def Sub(self, node): + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_y = self.graph.get_input_node(node, idx=1, copy=True) + inputs = { + "x": val_x, + "y": val_y, + } + attr = {"name": string(node.layer_name)} + node.fluid_code.add_layer("elementwise_sub", + inputs=inputs, + output=node, + param_attr=attr) + + def Pow(self, node): + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_y = self.graph.get_input_node(node, idx=1, copy=True) + inputs = { + "x": val_x, + "y": val_y, + } + attr = {"name": string(node.layer_name)} + node.fluid_code.add_layer("elementwise_pow", + inputs=inputs, + output=node, + param_attr=attr) + def Sum(self, node): val_inps = node.layer.input inputs = { - "x": val_inps[0], - "y": val_inps[1], + "x": self.graph.get_input_node(node, idx=0, copy=True), + "y": self.graph.get_input_node(node, idx=1, copy=True), } node.fluid_code.add_layer("elementwise_add", inputs=inputs, output=node) - for ipt in val_inps[2:]: + for idx, ipt in enumerate(val_inps[2:]): + y = self.graph.get_input_node(node, idx=idx, copy=True) inputs = { "x": node.layer_name, - "y": ipt, + "y": y, } node.fluid_code.add_layer("elementwise_add", inputs=inputs, output=node) def MatMul(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_y = self.graph.get_node(node.layer.input[1], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_y = self.graph.get_input_node(node, idx=1, copy=True) inputs = {"x": val_x, "y": val_y} attr = {"name": string(node.layer_name)} node.fluid_code.add_layer("matmul", @@ -667,11 +841,11 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def BatchNormalization(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_scale = self.graph.get_node(node.layer.input[1], copy=True) - val_b = self.graph.get_node(node.layer.input[2], copy=True) - val_mean = self.graph.get_node(node.layer.input[3], copy=True) - val_var = self.graph.get_node(node.layer.input[4], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_scale = self.graph.get_input_node(node, idx=1, copy=True) + val_b = self.graph.get_input_node(node, idx=2, copy=True) + val_mean = self.graph.get_input_node(node, idx=3, copy=True) + val_var = self.graph.get_input_node(node, idx=4, copy=True) self.omit_nodes.append(val_scale.layer_name) self.omit_nodes.append(val_b.layer_name) @@ -701,7 +875,7 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Transpose(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) perm = node.get_attr('perm') attr = {'perm': perm, "name": string(node.layer_name)} node.fluid_code.add_layer("transpose", @@ -710,12 +884,9 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Mul(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_y = self.graph.get_node(node.layer.input[1], copy=True) - - val_x_shape = val_x.out_shapes[0] + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_y = self.graph.get_input_node(node, idx=1, copy=True) val_y_shape = val_y.out_shapes[0] - slice_idx = 0 for dim in val_y_shape: if dim == 1: @@ -747,12 +918,9 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Div(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_y = self.graph.get_node(node.layer.input[1], copy=True) - - val_x_shape = val_x.out_shapes[0] + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_y = self.graph.get_input_node(node, idx=1, copy=True) val_y_shape = val_y.out_shapes[0] - slice_idx = 0 for dim in val_y_shape: if dim == 1: @@ -784,7 +952,7 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Relu(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) attr = {"name": string(node.layer_name)} node.fluid_code.add_layer("relu", inputs=val_x, @@ -792,8 +960,8 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def PRelu(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_slope = self.graph.get_node(node.layer.input[1], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_slope = self.graph.get_input_node(node, idx=1, copy=True) mode = 'channel' shape_slope = val_slope.out_shapes[0] @@ -811,20 +979,20 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Squeeze(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - squeeze_dims = node.get_attr('squeeze_dims') - attr = {'axes': squeeze_dims, "name": string(node.layer_name)} + val_x = self.graph.get_input_node(node, idx=0, copy=True) + axes = node.get_attr('axes') + attr = {'axes': axes, "name": string(node.layer_name)} node.fluid_code.add_layer("squeeze", inputs=val_x, output=node, param_attr=attr) def Identity(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) node.fluid_code.add_layer("assign", inputs=val_x, output=node) def MaxPool(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) auto_pad = node.get_attr('auto_pad', 'NOTSET') assert node.get_attr( @@ -839,10 +1007,10 @@ class ONNXOpMapper(OpMapper): fluid_op = 'pool{}d'.format(poolnd) assert 2 <= poolnd <= 3, 'only pool2d and pool3d is supported' - input_shape = val_x.out_shapes[0] paddings, val_x = self._pad_if_asymmetric(node, pads, val_x) if auto_pad == "SAME_UPPER" or auto_pad == "SAME_LOWER": + input_shape = val_x.out_shapes[0] pad_h = get_same_padding(input_shape[2], kernel_shape[0], strides[0]) pad_w = get_same_padding(input_shape[3], kernel_shape[1], @@ -864,7 +1032,7 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def GlobalAveragePool(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) val_y = self.graph.get_node(node.layer.output[0], copy=True) input_shape = val_x.out_shapes[0] output_shape = val_y.out_shapes[0] @@ -886,21 +1054,19 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def Conv(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_w = self.graph.get_node(node.layer.input[1], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_w = self.graph.get_input_node(node, idx=1, copy=True) val_y = self.graph.get_node(node.layer.output[0], copy=True) self.omit_nodes.append(val_w.layer_name) has_bias = len(node.layer.input) == 3 if has_bias: - val_b = self.graph.get_node(node.layer.input[2], copy=True) + val_b = self.graph.get_input_node(node, idx=2, copy=True) self.omit_nodes.append(val_b.layer_name) auto_pad = node.get_attr('auto_pad', 'NOTSET') - kernel_shape = val_w.out_shapes[0][2:] # OI... - assert kernel_shape == node.get_attr( - 'kernel_shape'), 'kernel_shape in attr unmatches value_info' # HW + kernel_shape = node.get_attr('kernel_shape') convnd = len(kernel_shape) assert 2 <= convnd <= 3, 'only conv2d and conv3d is supported' num_out_channels = val_w.out_shapes[0][0] # OI... @@ -941,9 +1107,9 @@ class ONNXOpMapper(OpMapper): param_attr=attr) def ConvTranspose(self, node): - val_x = self.graph.get_node(node.layer.input[0], copy=True) - val_w = self.graph.get_node(node.layer.input[1], copy=True) - val_b = self.graph.get_node(node.layer.input[2], copy=True) + val_x = self.graph.get_input_node(node, idx=0, copy=True) + val_w = self.graph.get_input_node(node, idx=1, copy=True) + val_b = self.graph.get_input_node(node, idx=2, copy=True) self.omit_nodes.append(val_w.layer_name) self.omit_nodes.append(val_b.layer_name) @@ -952,7 +1118,7 @@ class ONNXOpMapper(OpMapper): auto_pad = node.get_attr('auto_pad', 'NOTSET') out_padding = node.get_attr('output_padding', [0, 0]) - kernel_shape = node.get_attr('kernel_shape', val_w.out_shapes[0][2:]) + kernel_shape = node.get_attr('kernel_shape') assert kernel_shape, 'kernel_shape not inferred' convnd = len(kernel_shape) assert 2 <= convnd <= 3, 'only conv2d_transpose and conv3d_transpose supported' diff --git a/x2paddle/op_mapper/tf_op_mapper.py b/x2paddle/op_mapper/tf_op_mapper.py index 5453184169a68e881c56250048f4c6e23962710e..73b0da50d797e2eb9fca0c724499e0ed4fdbf367 100644 --- a/x2paddle/op_mapper/tf_op_mapper.py +++ b/x2paddle/op_mapper/tf_op_mapper.py @@ -785,6 +785,9 @@ class TFOpMapper(OpMapper): start = self.graph.get_node(node.layer.input[0], copy=True) limit = self.graph.get_node(node.layer.input[1], copy=True) delta = self.graph.get_node(node.layer.input[2], copy=True) + self.add_omit_nodes(start.layer_name, node.layer_name) + self.add_omit_nodes(limit.layer_name, node.layer_name) + self.add_omit_nodes(delta.layer_name, node.layer_name) if start.layer_type == "Const": start = start.value else: @@ -797,9 +800,6 @@ class TFOpMapper(OpMapper): delta = delta.value else: delta = self.decoder.infer_tensor(delta) - self.add_omit_nodes(start.layer_name, node.layer_name) - self.add_omit_nodes(limit.layer_name, node.layer_name) - self.add_omit_nodes(delta.layer_name, node.layer_name) inputs = {"start": start, "end": limit, "step": delta} attr = {"dtype": string(node.dtype)} diff --git a/x2paddle/op_mapper/tf_op_mapper_nhwc.py b/x2paddle/op_mapper/tf_op_mapper_nhwc.py index 9399c6d7a13f26fee17af9a76a52f6d5de5dba8b..cc1e355c866b90fb3cdc8e0d2e7065e6b357c12d 100644 --- a/x2paddle/op_mapper/tf_op_mapper_nhwc.py +++ b/x2paddle/op_mapper/tf_op_mapper_nhwc.py @@ -74,7 +74,7 @@ class TFOpMapperNHWC(OpMapper): unsupported_ops = set() sys.stderr.write("Total nodes: {}\n".format(len(self.graph.topo_sort))) for i, node_name in enumerate(self.graph.topo_sort): - sys.stderr.write("\rConverting node {} ... ".format(i)) + sys.stderr.write("\rConverting node {} ... ".format(i + 1)) node = self.graph.get_node(node_name) op = node.layer_type if op in self.directly_map_ops: @@ -99,7 +99,7 @@ class TFOpMapperNHWC(OpMapper): for op in unsupported_ops: print("========== {} ============".format(op)) sys.exit(-1) - sys.stderr.write("\nDone\n") + sys.stderr.write("\nDone!\n") def add_omit_nodes(self, in_node_name, out_node_name): in_node = self.graph.get_node(in_node_name)