""" onnx modifier: provide a conviennt way to modify onnx model 1. add node 2. remove node 3. modify node 4. query node """ from collections import defaultdict, deque import os import os.path as osp import shutil import tempfile from typing import List, Dict, Set, Tuple, Optional, Union import uuid import warnings import numpy as np import onnx from onnx import AttributeProto, numpy_helper from onnx import shape_inference from onnx.helper import make_attribute, make_node, make_opsetid, make_tensor, \ tensor_dtype_to_np_dtype from onnxconverter_common import float16 import tqdm # from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE SUPPORT_DTYPES = [ 'BOOL', 'STRING', 'BFLOAT16', 'DOUBLE', 'FLOAT', 'FLOAT16', 'INT16', 'INT32', 'INT4', 'INT64', 'INT8', 'UINT16', 'UINT32', 'UINT4', 'UINT64', 'UINT8', ] SUPPORT_DTYPES.extend([dt.lower() for dt in SUPPORT_DTYPES]) class Node: def __init__(self, onnx_modifier=None, obj=None, index=None): self.onnx_modifier = onnx_modifier self.obj = obj self.index = index @property def name(self): return self.obj.name @property def op_type(self): return self.obj.op_type @property def inputs(self): return self.obj.input @property def outputs(self): return self.obj.output @property def input_names(self): return self.inputs @property def output_names(self): return self.outputs def check_modifier(self): if self.onnx_modifier is None: raise RuntimeError("onnx_modifier is not initialized") @property def prev_nodes(self): self.check_modifier() return self.onnx_modifier.get_prev_nodes(self) @property def next_nodes(self): self.check_modifier() return self.onnx_modifier.get_next_nodes(self) def replace_input(self, old_name, new_name): assert old_name in self.obj.input, \ f'"{old_name}" not in input name list of node named "{self.name}"' for i, in_name in enumerate(self.obj.input): if in_name == old_name: self.set_input(i, new_name) def set_input(self, index, name): # assert index < len(self.obj.input), "index out of range" # orig_name = self.obj.input[index] # self.obj.input[index] = name assert index < len(self.onnx_modifier.graph.node[self.index].input), "index out of range" orig_name = self.onnx_modifier.graph.node[self.index].input[index] self.onnx_modifier.graph.node[self.index].input[index] = name self.check_modifier() # Can not execute connection.pop_to_node() method directly. # When node inputs contain multiple orig_name, need to remain the node in to_nodes. if list(self.onnx_modifier.graph.node[self.index].input).count(orig_name) == 0: self.onnx_modifier.connection_map[orig_name].pop_to_node(self) if name not in self.onnx_modifier.connection_map: self.onnx_modifier.connection_map[name] = Connection(name, self.onnx_modifier) self.onnx_modifier.connection_map[name].add_to_node(self) def set_inputs(self, names): assert len(names) == len(self.obj.input), "number of inputs does not match" assert all(isinstance(name, str) for name in names), "input names must be strings" self.obj.input[:] = names def set_output(self, index, name): assert index < len(self.obj.output), "index out of range" orig_name = self.obj.output[index] self.obj.output[index] = name self.check_modifier() self.onnx_modifier.connection_map[orig_name].clear_from_node() if name not in self.onnx_modifier.connection_map: self.onnx_modifier.connection_map[name] = Connection(name, self.onnx_modifier) self.onnx_modifier.connection_map[name].set_from_node(self) def set_outputs(self, names): assert len(names) == len(self.obj.output), "number of outputs does not match" assert all(isinstance(name, str) for name in names), "output names must be strings" self.obj.output[:] = names @property def attrs(self): attrs = {} for attr in self.obj.attribute: if attr.type == AttributeProto.FLOAT: # 1 value = attr.f elif attr.type == AttributeProto.INT: # 2 value = attr.i elif attr.type == AttributeProto.STRING: # 3 value = attr.s.decode('utf-8') elif attr.type == AttributeProto.TENSOR: # 4 value = numpy_helper.to_array(attr.t) elif attr.type == AttributeProto.FLOATS: # 6 value = list(attr.floats) elif attr.type == AttributeProto.INTS: # 7 value = list(attr.ints) else: value = f"Unsupported type: {attr.type}" attrs[attr.name] = value return attrs def set_attribute(self, name, value, name2attr=None): if not name2attr: name2attr = {} for attr in self.obj.attribute: name2attr[attr.name] = attr if name in name2attr: if isinstance(value, float): name2attr[name].f = value name2attr[name].type = AttributeProto.FLOAT elif isinstance(value, int): name2attr[name].i = value name2attr[name].type = AttributeProto.INT elif isinstance(value, str): name2attr[name].s = value.encode('utf-8') name2attr[name].type = AttributeProto.STRING elif isinstance(value, np.ndarray): name2attr[name].ClearField("t") name2attr[name].t.CopyFrom(numpy_helper.from_array(value)) elif isinstance(value, list): is_all_float = all(isinstance(x, float) for x in value) is_all_int = all(isinstance(x, int) for x in value) assert is_all_float or is_all_int if is_all_float: name2attr[name].ClearField("floats") name2attr[name].floats.extend(value) name2attr[name].type = AttributeProto.FLOATS else: name2attr[name].ClearField("ints") name2attr[name].ints.extend(value) name2attr[name].type = AttributeProto.INTS else: if isinstance(value, np.ndarray): value = numpy_helper.from_array(value) self.obj.attribute.append(make_attribute(name, value)) def set_attributes(self, attr_dict): name2attr = {} for attr in self.obj.attribute: name2attr[attr.name] = attr for name, value in attr_dict.items(): self.set_attribute(name, value, name2attr) class Connection: def __init__(self, conn_name, onnx_modifier=None): self.name = conn_name self.onnx_modifier = onnx_modifier self.from_node = None self.to_nodes = [] self.to_node_names = set() def check_modifier(self): if self.onnx_modifier is None: raise RuntimeError("onnx_modifier is not initialized") def set_from_node(self, node: str | Node): if isinstance(node, str): self.check_modifier() _node = self.onnx_modifier.get_node(Node) assert node is not None, f'No node named "{node}" in onnx graph!' elif isinstance(node, Node): _node = node else: raise TypeError(f"Connection.set_from_node except input argument type" \ f" is str or Node, but received: {type(node)}") self.from_node = _node def clear_from_node(self): self.from_node = None def add_to_node(self, node: str | Node): if isinstance(node, str): _name = node self.check_modifier() _node = self.onnx_modifier.get_node(Node) assert node is not None, f'No node named "{node}" in onnx graph!' elif isinstance(node, Node): _name = node.name _node = node else: raise TypeError(f"Connection.add_to_node except input argument type" \ f" is str or Node, but received: {type(node)}") if _name not in self.to_node_names: self.to_node_names.add(_name) self.to_nodes.append(_node) def pop_to_node(self, node: str | Node): if isinstance(node, str): _name = node self.check_modifier() _node = self.onnx_modifier.get_node(Node) assert node is not None, f'No node named "{node}" in onnx graph!' elif isinstance(node, Node): _name = node.name _node = node else: raise TypeError(f"Connection.pop_to_node except input argument type" \ f" is str or Node, but received: {type(node)}") if _name not in self.to_node_names: raise ValueError(f'Node "{_name}" not in target nodes of connction "{self.name}"!') self.to_node_names.remove(_name) for i in range(len(self.to_nodes)): if self.to_nodes[i].name == _name: return self.to_nodes.pop(i) else: raise RuntimeError("to_nodes dismatch to_node_names!") class ONNXModifier: def __init__(self, onnx_path): self.onnx_path = onnx_path self.node_map = {} self.initializer_map = {} self.sparse_initializer_map = {} self.connection_map = {} self.value_info_map = {} self.parse_onnx(self.onnx_path) def parse_onnx(self, onnx_path): model = onnx.load(onnx_path) self.model = model self.domain = model.domain self.graph = model.graph self.ir_version = model.ir_version self.mdoel_version = model.model_version self.opset_import = model.opset_import self.update_map() def add_node_name_if_nameless(self, node: Node): if not hasattr(self, "node_suffixes"): self.name_suffixes = set() if node.name == "" or node.name == None: suffix = None while True: suffix = uuid.uuid4().hex[:8] if suffix not in self.name_suffixes: break node.obj.name = node.op_type + "_" + suffix def add_opset_import(self, domain: str, version: int): self.model.opset_import.append(make_opsetid(domain, version)) @property def input_names(self): return [i.name for i in self.graph.input] @property def output_names(self): return [o.name for o in self.graph.output] def add_input(self, name, dtype='float32', shape=None): assert dtype in set(SUPPORT_DTYPES) self.create_value_info(name, dtype=dtype, shape=shape) new_input = self.value_info_map.pop(name) _new_input = self.graph.value_info.pop() assert id(new_input) == id(_new_input) assert name == new_input.name self.graph.input.append(new_input) return new_input def add_output(self, name, new_name=None, shape=None): if name not in self.value_info_map: raise ValueError(f"{name} not in onnx_modifier.value_info_map") index = None for i, v in enumerate(self.graph.value_info): if v.name == name: index = i break else: raise ValueError(f"{name} not in model.graph.value_info") value_info = self.value_info_map.pop(name) assert value_info.name == name assert id(value_info) == id(self.graph.value_info[index]) self.graph.value_info.pop(index) if shape is not None: tensor_type = onnx.helper.make_tensor_type_proto( elem_type=value_info.type.tensor_type.elem_type, shape=shape ) value_info.type.CopyFrom(tensor_type) if new_name is None: self.graph.output.append(value_info) else: from_node = self.get_from_node(name) to_nodes = self.get_to_nodes(name) for i, output_name in enumerate(from_node.output_names): if output_name == name: from_node.set_output(i, new_name) for node in to_nodes: node.replace_input(name, new_name) value_info.name = new_name self.graph.output.append(value_info) def remove_output(self, name): """根据名称删除模型输出""" assert name in self.output_names # print("need remove output name:", name) index = None for i, out in enumerate(self.graph.output): # print(f"current(index={i}) output name:", out.name) if out.name == name: index = i break else: raise RuntimeError(f"ONNX graphx not has a output named '{name}'.") self.graph.output.pop(index) def get_node(self, name_or_index: Union[str, int]): """根据节点名称或索引获取节点实例""" if isinstance(name_or_index, str): if name_or_index in self.node_map: return self.node_map.get(name_or_index, None) elif isinstance(name_or_index, int): if name_or_index < len(self.graph.node): return self.node_map.get( self.graph.node[name_or_index].name, None) else: raise ValueError(f"Node index {name_or_index} out of range") def get_nodes(self, *op_types: str): """根据节点类型获取节点实例""" assert len(op_types) >= 1 op_types_set = set(op_types) node_names = [node.name for node in self.graph.node if node.op_type in op_types_set] nodes = [self.node_map[name] for name in node_names] return nodes def get_initializer(self, name: str): """根据initializer名称获取initializer""" return self.initializer_map.get(name) def get_connection(self, name: str): """根据边名称获取边""" return self.connection_map.get(name) def get_from_node(self, conn: Union[str, Connection]): """获取某条边的输入节点名""" if isinstance(conn, str): assert conn in self.connection_map, f"Connection {conn} not in connection_map!" return self.connection_map[conn].from_node elif isinstance(conn, Connection): return conn.from_node else: raise TypeError(f"Invalid connection type {type(conn)}") def get_to_nodes(self, conn: Union[str, Connection]): """获取某条边的输出节点""" if isinstance(conn, str): assert conn in self.connection_map, f"Connection {conn} not in connection_map!" return self.connection_map[conn].to_nodes elif isinstance(conn, Connection): return conn.to_nodes else: raise TypeError(f"Invalid connection type {type(conn)}") def get_prev_nodes(self, node: Union[str, Node]): """获取某节点的上游输入节点""" if isinstance(node, str): node = self.node_map[node] elif isinstance(node, Node): pass else: raise TypeError(f"Invalid node type {type(node)}") nodes = [] for conn_name in node.inputs: from_node = self.get_from_node(conn_name) if from_node: nodes.append(from_node) return nodes def get_next_nodes(self, node: Union[str, Node]): """获取某节点的下游节点""" if isinstance(node, str): node = self.node_map[node] elif isinstance(node, Node): pass else: raise TypeError(f"Invalid node type {type(node)}") nodes = [] for conn_name in node.outputs: to_nodes = self.get_to_nodes(conn_name) nodes.extend(to_nodes) return nodes def create_node(self, op_type, op_name, inputs, outputs, doc_string=None, domain=None, index=None, **attrs): """创建一个新节点""" onnx_node = make_node(op_type, inputs, outputs, name=op_name, doc_string=doc_string, domain=domain, **attrs) if index is None: self.graph.node.append(onnx_node) index = len(self.graph.node) - 1 else: assert index <= len(self.graph.node), "index out of range" self.graph.node.insert(index, onnx_node) for i in range(index + 1, len(self.graph.node)): node_name = self.graph.node[i].name old_idx = self.node_map[node_name].index assert old_idx == i - 1, \ f"Node {node_name} index conflict: {old_idx} != {i - 1}" self.node_map[node_name].index = i new_node = Node(self, self.graph.node[index], index) self.node_map[op_name] = new_node for in_name in new_node.input_names: if in_name not in self.value_info_map: self.create_value_info(in_name, dtype="float") if in_name not in self.connection_map: self.connection_map[in_name] = Connection(in_name, self) self.connection_map[in_name].add_to_node(new_node) for out_name in new_node.output_names: if out_name not in self.value_info_map: self.create_value_info(out_name, dtype="float") if out_name not in self.connection_map: self.connection_map[out_name] = Connection(out_name, self) self.connection_map[out_name].set_from_node(new_node) return new_node def create_initializer(self, name, value: np.ndarray): """创建一个 initializer""" assert name not in self.initializer_map, f"initializer {name} already exists!" init_node = numpy_helper.from_array(value, name=name) use_external_data = value.nbytes / 1024 / 1024 / 1024 > 2 if use_external_data: print("use external data:", name) init_node.data_location = onnx.TensorProto.EXTERNAL location = name.replace('/', '+') + '.data' onnx.external_data_helper.set_external_data(init_node, location) with tempfile.TemporaryDirectory() as tmp_dir: onnx.external_data_helper.save_external_data(init_node, tmp_dir) init_node.ClearField("raw_data") self.graph.initializer.append(init_node) onnx.external_data_helper.load_external_data_for_tensor( self.graph.initializer[-1], tmp_dir) del self.graph.initializer[-1].external_data[:] self.graph.initializer[-1].ClearField("data_location") else: self.graph.initializer.append(init_node) self.initializer_map[name] = self.graph.initializer[-1] return self.graph.initializer[-1] def create_value_info(self, name, dtype=None, shape=None): if dtype is None: elem_type = None else: assert isinstance(dtype, str) assert dtype in set(SUPPORT_DTYPES) elem_type = getattr(onnx.TensorProto, dtype.upper()) value_info = onnx.helper.make_tensor_value_info(name=name, elem_type=elem_type, shape=shape) self.graph.value_info.append(value_info) self.value_info_map[name] = self.graph.value_info[-1] return self.graph.value_info[-1] def get_initializer_value(self, name): """获取initializer的数值""" init = self.get_initializer(name) return numpy_helper.to_array(init) def set_initializer_value(self, name, value: np.ndarray): """为initializer设置新的数值""" init = self.get_initializer(name) # 检查形状和类型 old_shape = list(init.dims) new_shape = list(value.shape) # old_dtype = TENSOR_TYPE_TO_NP_TYPE.get(init.data_type, None) old_dtype = tensor_dtype_to_np_dtype(init.data_type) new_dtype = value.dtype if old_shape != new_shape: warn_message = f"Initailizer {name} shape changed: {old_shape} -> {new_shape}" warnings.warn(warn_message, RuntimeWarning) if old_dtype is not None and old_dtype != new_dtype: warn_message = f"Initailizer {name} dtype changed: {old_dtype} -> {new_dtype}" warnings.warn(warn_message, RuntimeWarning) new_tensor_proto = numpy_helper.from_array(value, name=name) init.CopyFrom(new_tensor_proto) def connect_node(self, node, inputs_map, outputs_map): """将某个节点与其上下游节点连接起来 Args: node: Node inputs_map: [(node0, out_idx0), (node1, out_idx1), ...] outputs_map: [(node0, in_idx0), (node1, in_idx1), ...] """ # 在连接 A -> B 时，若 A 的输出名与 B 的输入名冲突时，优先使用 A 的输出名， # 即：B.input[i] = A.output[j] for i, (n, j) in enumerate(inputs_map): if isinstance(n, str): n = self.node_map[n] assert j < len(n.outputs), \ f"output index {i} out of node {n.name} outputs range" node.set_input(i, n.outputs[j]) for name, (n, i) in zip(node.outputs, outputs_map): if isinstance(n, str): n = self.node_map[n] assert i < len(n.outputs), \ f"output index {i} out of node {n.name} outputs range" n.set_output(i, name) # TODO: update self.connection_map def pop_node(self, node: Union[str, Node, int], auto_connect=True): """根据节点名称或索引移除节点""" if isinstance(node, str): node = self.node_map.get(node, None) if node is None: return None index = node.index assert node.name == self.graph.node[index].name elif isinstance(node, int): if node >= len(self.graph.node): raise ValueError(f"Node index {node} out of range") index = node elif isinstance(node, Node): index = node.index else: raise ValueError(f"Invalid node name or index: {node}") for i in range(index + 1, len(self.graph.node)): node = self.graph.node[i] self.node_map[node.name].index -= 1 # print(f"node_name={self.graph.node[index].name} node_index={index}") _node_obj = self.graph.node[index] _node = self.get_node(_node_obj.name) next_nodes = self.get_next_nodes(_node) self.graph.node.pop(index) self.node_map.pop(_node_obj.name) # automatic connecting edges if auto_connect and len(_node.inputs) == 1 and len(_node.outputs) == 1: # self.connection_map[_node.inputs[0]].pop_to_node(_node) for next_node in next_nodes: next_node.replace_input( _node.outputs[0], _node.inputs[0]) # self.connection_map[_node.inputs[0]].add_to_node(next_node) # self.connection_map.pop(_node.outputs[0]) # update connection_map for in_name in _node.input_names: if _node.name in self.connection_map[in_name].to_node_names: self.connection_map[in_name].pop_to_node(_node) for i, out_name in enumerate(_node.output_names): self.connection_map[out_name].clear_from_node() return _node def remove_nodes(self, nodes: List[str | Node], auto_connect=False): """同时删除多个节点""" indices = set() _nodes = [] invalid_nodes = set() for node in nodes: if isinstance(node, str): if node in self.node_map: node = self.node_map[node] if node.index not in indices: _nodes.append(node) indices.add(node.index) else: invalid_nodes.add(node) elif isinstance(node, Node): if node.index not in indices: _nodes.append(node) indices.add(node.index) else: invalid_nodes.add(node) _nodes.sort(key=lambda x:x.index, reverse=True) use_progress_bar = len(_nodes) > 500 if use_progress_bar: pbar = tqdm.tqdm(total=len(_nodes), desc="Removing nodes") for node in _nodes: self.pop_node(node, auto_connect=auto_connect) if use_progress_bar: pbar.update(1) if use_progress_bar: pbar.close() # print(f"{len(nodes) - len(invalid_nodes)} nodes have been removed.") # if len(invalid_nodes) > 0: # print(f"find {len(invalid_nodes)} invalid nodes:\n", invalid_nodes) def pop_initializer(self, init_name: str, update_node_inputs: bool = True): """根据initializer名字移除initializer""" _init1 = self.initializer_map.pop(init_name) init_index = None for i in range(len(self.graph.initializer)): if self.graph.initializer.name == init_name: init_index = i break else: raise ValueError(f"Not existing a Initializer named {init_name}") _init2 = self.graph.initializer.pop(init_index) assert id(_init1) == id(_init2) # if update_node_inputs and init_name in self.connection_map: # to_nodes = self.get_to_nodes(init_name) # self.connection_map.pop(init_name) # for node in to_nodes: # num_inputs = len(node.inputs) # for i in range(num_inputs-1, -1, -1): # if node.inputs[i] == init_name: # node.inputs.pop(i) return _init1 def update_map(self): """更新connection_map与node_map""" self.node_map.clear() self.connection_map.clear() self.initializer_map.clear() self.sparse_initializer_map.clear() self.value_info_map.clear() for i, node in enumerate(self.graph.node): new_node = Node(self, node, i) self.add_node_name_if_nameless(new_node) self.node_map[node.name] = new_node for conn_name in node.input: if conn_name not in self.connection_map: self.connection_map[conn_name] = Connection(conn_name, self) self.connection_map[conn_name].add_to_node(new_node) for conn_name in node.output: if conn_name not in self.connection_map: self.connection_map[conn_name] = Connection(conn_name, self) self.connection_map[conn_name].set_from_node(new_node) for i, node in enumerate(self.graph.initializer): self.initializer_map[node.name] = node for i, node in enumerate(self.graph.sparse_initializer): self.sparse_initializer_map[node.name] = [node, i] for i, conn in enumerate(self.graph.value_info): self.value_info_map[conn.name] = conn def find_unuseful_nodes(self): """寻找没有用到的节点""" end_names = set() for output_name in self.output_names: end_names.add(self.get_from_node(output_name).name) unuseful_names = set() for node in self.node_map.values(): if node.name in end_names: continue next_nodes = self.get_next_nodes(node) if len(next_nodes) == 0: unuseful_names.add(node.name) model_output_names = set(self.output_names) q = deque([self.node_map[name] for name in unuseful_names]) while len(q) != 0: node = q.popleft() prev_nodes = self.get_prev_nodes(node) for node1 in prev_nodes: next_nodes = self.get_next_nodes(node1) next_names = set([node2.name for node2 in next_nodes]) # if (next_names - end_names).issubset(unuseful_names): if next_names.issubset(unuseful_names): if node1.name not in unuseful_names and set(node1.output_names).isdisjoint(model_output_names): q.append(node1) unuseful_names.add(node1.name) unuseful_nodes = [self.node_map[name] for name in unuseful_names] return unuseful_nodes def remove_trash(self): """ 1. 移除无用的节点 2. 移除无用的initializer 3. 移除没有输入节点的connection 4. 移除没有用到的模型输入与输出 5. 移除没有用到的value_info """ self.update_map() unuseful_nodes = self.find_unuseful_nodes() print(f"Find unuseful {len(unuseful_nodes)} nodes!") for i, node in enumerate(unuseful_nodes): print(f"remove unuseful node {i}:", node.name) self.remove_nodes(unuseful_nodes) self.update_map() all_node_inputs = set() for node in self.node_map.values(): all_node_inputs.update(node.input_names) # remove unuseful initializers cnt = 0 for init_name in list(self.initializer_map.keys()): if init_name in all_node_inputs: continue index = None for i, init in enumerate(self.graph.initializer): if init.name == init_name: index = i break else: raise ValueError( f"{init_name} not in model.graph.initializer") print(f"remove unuseful initializer {cnt}:", init_name) self.graph.initializer.pop(index) cnt += 1 # remove unuseful sparse_initializers cnt = 0 for init_name in list(self.sparse_initializer_map.keys()): if init_name in all_node_inputs: continue index = None for i, init in enumerate(self.graph.sparse_initializer): if init.name == init_name: index = i break else: raise ValueError( f"{init_name} not in model.graph.sparse_initializer") print(f"remove unuseful sparse initializer {cnt}:", init_name) self.graph.sparse_initializer.pop(index) cnt += 1 self.update_map() # remove unuseful inputs and outputs for in_name in self.input_names: # print(in_name, [n.name for n in self.get_to_nodes(in_name)]) if len(self.get_to_nodes(in_name)) != 0: continue for i, _in in enumerate(self.graph.input): if in_name == _in.name: self.graph.input.pop(i) break for out_name in self.output_names: # print(out_name, self.get_from_node(out_name).name) if self.get_from_node(out_name) is not None: continue for i, _out in enumerate(self.graph.output): if out_name == _out.name: self.graph.output.pop(i) break self.update_map() # remove unuseful value_info cnt = 0 num_value_info = len(self.graph.value_info) for i in range(num_value_info-1, -1, -1): v = self.graph.value_info[i] if v.name not in self.connection_map: self.graph.value_info.pop(i) print(f"remove unuseful value_info {cnt}:", v.name) cnt += 1 self.update_map() def infer_shape(self, strict_mode=False): for vi in self.graph.value_info: if vi.type.HasField("tensor_type"): vi.type.tensor_type.ClearField("shape") model = shape_inference.infer_shapes(self.model, strict_mode=strict_mode) self.model = model self.domain = model.domain self.graph = model.graph self.ir_version = model.ir_version self.mdoel_version = model.model_version self.opset_import = model.opset_import self.update_map() def infer_node_shpe(self, node): input_shapes = [] input_dtypes = [] for input_name in node.inputs: value_info = self.value_info_map[input_name] input_shapes.append(value_info.type.tensor_type.dims) input_dtypes.append(value_info.type.tensor_type.type) shape_inference.infer_node_outputs(node.obj, input_shapes, input_dtypes) def convert_float_to_float16(self): self.model = float16.convert_float_to_float16(self.model, keep_io_types=True) def save(self, save_path, save_as_external_data=False, all_tensors_to_one_file=True): self.remove_trash() external_data_name = osp.basename(save_path) + '.data' external_data_path = osp.join(osp.dirname(save_path), external_data_name) if save_as_external_data and osp.isfile(external_data_path): os.remove(external_data_path) onnx.save(self.model, save_path, save_as_external_data=save_as_external_data, all_tensors_to_one_file=all_tensors_to_one_file, location=external_data_name, size_threshold=1024, convert_attribute=False)