def _type_to_proto(t: torch._C.TensorType) -> onnx.TypeProto: if t.kind() == "NoneType": return onnx.TypeProto() ret: onnx.TypeProto = onnx.TypeProto() ret.denotation = repr(t) if t.kind() == "ListType": ret.sequence_type.elem_type.CopyFrom(_type_to_proto(cast(torch._C.TensorType, t.getElementType()))) return ret if t.kind() == "IntType": ret.tensor_type.elem_type = onnx.TensorProto.DataType.INT64 ret.tensor_type.shape.CopyFrom(onnx.TensorShapeProto()) return ret assert t.kind() == "TensorType", f"Not Tensor type(actual: {t.kind()}): {t}" if t.scalarType() is None: ret.tensor_type.elem_type = onnx.TensorProto.DataType.UNDEFINED else: ret.tensor_type.elem_type = int( # type: ignore sym_hel.cast_pytorch_to_onnx[t.scalarType()] # type: ignore[index] ) ret.tensor_type.shape.CopyFrom(onnx.TensorShapeProto()) if t.sizes() is not None: for s in t.sizes(): # type: ignore d = ret.tensor_type.shape.dim.add() d.dim_value = s assert ret.tensor_type.HasField("shape") return ret
def dump_test_inputs_outputs(inputs, outputs, test_data_dir): if not os.path.exists(test_data_dir): os.makedirs(test_data_dir) for typ, values in [('input', inputs), ('output', outputs)]: for i, (value_info, value) in enumerate(values): name = value_info.name if isinstance(value, list): assert value digits = len(str(len(value))) for j, v in enumerate(value): filename = os.path.join( test_data_dir, '%s_%d_%s.pb' % (typ, i, str(j).zfill(digits))) tensor = tensor_from_array(v, name) with open(filename, 'wb') as f: f.write(tensor.SerializeToString()) value_info.type.CopyFrom(onnx.TypeProto()) sequence_type = value_info.type.sequence_type tensor_type = sequence_type.elem_type.tensor_type tensor_type.elem_type = tensor.data_type else: filename = os.path.join(test_data_dir, '%s_%d.pb' % (typ, i)) if value is None: if get_test_args().allow_unused_params: continue raise RuntimeError('Unused parameter: %s' % name) tensor = tensor_from_array(value, name) with open(filename, 'wb') as f: f.write(tensor.SerializeToString()) vi = onnx.helper.make_tensor_value_info( name, tensor.data_type, tensor.dims) value_info.CopyFrom(vi)
def add_initializers_into_inputs(model: onnx.ModelProto) -> onnx.ModelProto: for x in model.graph.initializer: input_names = [x.name for x in model.graph.input] if x.name not in input_names: shape = onnx.TensorShapeProto() for dim in x.dims: shape.dim.extend([onnx.TensorShapeProto.Dimension(dim_value=dim)]) model.graph.input.extend( [onnx.ValueInfoProto(name=x.name, type=onnx.TypeProto(tensor_type=onnx.TypeProto.Tensor(elem_type=x.data_type, shape=shape)))]) return model
def add_initializers_into_inputs(model: onnx.ModelProto) -> onnx.ModelProto: # Due to a onnx bug, https://github.com/onnx/onnx/issues/2417, we need to add missing initializers into inputs for x in model.graph.initializer: input_names = [x.name for x in model.graph.input] if x.name not in input_names: shape = onnx.TensorShapeProto() for dim in x.dims: shape.dim.extend( [onnx.TensorShapeProto.Dimension(dim_value=dim)]) model.graph.input.extend([ onnx.ValueInfoProto( name=x.name, type=onnx.TypeProto(tensor_type=onnx.TypeProto.Tensor( elem_type=x.data_type, shape=shape))) ]) return model
def stan_select_model_inputs_outputs(model, dtype, inputs, outputs, io_shapes): """ a modificiation of select_model_input_outputs from sklearn-on Takes a model and changes its inputs and outputs :param model: *ONNX* model :param inputs: new inputs :return: modified model The function removes unneeded nodes. """ if dtype == np.float32: elem_type = onnx.TensorProto.FLOAT else: assert dtype == np.float64 elem_type = onnx.TensorProto.DOUBLE if inputs is None: raise NotImplementedError("Parameter inputs cannot be empty.") if outputs is None: raise NotImplementedError("Parameter inputs cannot be empty.") if not isinstance(inputs, list): inputs = [inputs] if not isinstance(outputs, list): outputs = [outputs] ########## mark_var = { } # keys are (input or node output) names, vals 1 = keep, 0 = delete for out in enumerate_model_node_outputs(model): mark_var[out] = 0 for inp in model.graph.input: mark_var[inp.name] = 0 for out in outputs: if out not in mark_var: raise ValueError( "Desired Output '{}' not found in model.".format(out)) initializers = [i.name for i in model.graph.initializer] for inp in inputs: if inp not in mark_var: raise ValueError( "Desired Input '{}' not found in model.".format(inp)) if inp not in initializers: mark_var[inp] = 1 nodes = list(enumerate(model.graph.node)) mark_op = { } # these are the marks for the node indices, 1 = keep, 0 = delete for node in nodes: mark_op[node[0]] = 0 # We mark all the nodes we need to keep. nb = 1 # number marked... used as a termination condition keep_initializers = [] while nb > 0: nb = 0 for index, node in nodes: if mark_op[index] == 1: # node was already processed, skip continue mod = False # is this a newly-marked node? node_initializers = [] for inp in node.input: if inp in outputs: continue if not inp in mark_var or mark_var.get(inp, 0) == 0: node_initializers.append(inp) # was initializer elif mark_var[inp] == 1: # make the node because its input was marked mark_op[index] = 1 mod = True for out in node.output: if out in inputs: continue if mark_var[out] == 1: # mark the node because the output was marked mark_op[index] = 1 mod = True if not mod: # none of the node's inputs were marked, skip it continue keep_initializers += node_initializers nb += 1 # mark the node and all its inputs / outputs for out in node.output: if mark_var.get(out, 0) == 1: continue if out in outputs: continue mark_var[out] = 1 nb += 1 for inp in node.input: if mark_var.get(inp, 0) == 1: continue if inp in inputs: continue mark_var[inp] = 1 nb += 1 # All nodes verifies mark_op[node.name] == 1 keep_nodes = [node[1] for node in nodes if mark_op[node[0]] == 1] var_in = [] for inp in inputs: nt = onnx.TypeProto() nt.tensor_type.elem_type = elem_type # inputs need shape info, which is not in the graph! shape = io_shapes[inp] for s in shape: nt.tensor_type.shape.dim.add() nt.tensor_type.shape.dim[-1].dim_value = s value_info = ValueInfoProto(type=nt) value_info.name = inp var_in.append(value_info) # add initializers to inputs for i in model.graph.input: if i.name in keep_initializers: var_in.append(i) var_out = [] for out in outputs: nt = onnx.TypeProto() nt.tensor_type.elem_type = elem_type # inputs need shape info, which is not in the graph! shape = io_shapes[out] for s in shape: nt.tensor_type.shape.dim.add() nt.tensor_type.shape.dim[-1].dim_value = s value_info = ValueInfoProto(type=nt) value_info.name = out var_out.append(value_info) init_out = [ init for init in model.graph.initializer if init.name in keep_initializers ] graph = make_graph(keep_nodes, model.graph.name, var_in, var_out, init_out) #print(f"making model with inputs {inputs} / outputs {outputs} and nodes len: {len(keep_nodes)}") onnx_model = make_model_with_graph(model, graph) return onnx_model
def generate_onnx(self) -> onnx.ModelProto: # Convert prim and aten nodes to ONNX by using symbolic functions self.original_g: torch._C.Graph = self.g.copy() target_nodes = list(self.g.nodes()) for n in target_nodes: self.generate_onnx_node(self.g, n) # Remove old prim and aten nodes by running DCE # After nodes is emited to ONNX nodes, all side effects should be removed self.run_jit_pass( torch._C._jit_pass_dce_allow_deleting_nodes_with_side_effects, self.g # type: ignore[attr-defined] ) # Run again to remove nodes only depending to aten node self.run_jit_pass( torch._C._jit_pass_dce_allow_deleting_nodes_with_side_effects, self.g # type: ignore[attr-defined] ) # TODO(twata): Remove unnecessary outputs. Graph#eraseOutput isn't available # while self.g.outputsSize() > len(self.outputs): # self.g.eraseOutput(self.g.outputsSize() - 1) self.optimize_onnx(self.g) self.log("ONNX graph", self.g) onnx_nodes, onnx_vars, val_tab = self.generate_proto_nodes(self.g, {}, {}) def onnx_value(v: torch._C.Value, name: ONNXValueID) -> onnx.ValueInfoProto: return onnx.helper.make_value_info( name, None if v.type() is None else _type_to_proto(cast(torch._C.TensorType, v.type())), doc_string=None if self.strip_doc_string else repr(v), ) def apply_dynamic_axes_info(out: onnx.ValueInfoProto, k: str) -> None: assert isinstance(self.dynamic_axes, dict) info = self.dynamic_axes.get(k, None) if info is None: return None if isinstance(info, list): ret: Dict[int, str] = {} for idx, axis in enumerate(info): ret[axis] = f"{k}_dynamic_axes_{idx + 1}" info = ret for axis, name in info.items(): out.type.tensor_type.shape.dim[axis].ClearField("dim_value") out.type.tensor_type.shape.dim[axis].dim_param = name # Values onnx_inputs: List[onnx.ValueInfoProto] = [] inout_names: List[str] = [] self_count = 0 for idx, v in enumerate(self.g.inputs()): if self.is_self(v): # Skip module's self input self_count += 1 continue if len(v.uses()) == 0: warnings.warn(f"Unused input: {v}") continue k = val_tab[_unique_id(v)] inout_names.append(k) onnx_inputs.append(onnx_value(v, k)) _apply_tensor_info_to_value_info(onnx_inputs[-1], self.inputs[idx - self_count]) apply_dynamic_axes_info(onnx_inputs[-1], k) if self.keep_initializers_as_inputs: for _, t_p in onnx_vars.items(): i_t: onnx.TypeProto = onnx.TypeProto() i_t.tensor_type.elem_type = t_p.data_type for d in t_p.dims: d_p = onnx.TensorShapeProto.Dimension() d_p.dim_value = d i_t.tensor_type.shape.dim.append(d_p) onnx_inputs.append(onnx.helper.make_value_info( t_p.name, i_t, )) onnx_outputs: List[onnx.ValueInfoProto] = [] for idx, v in enumerate(self.g.outputs()): k = val_tab[_unique_id(v)] inout_names.append(k) onnx_outputs.append(onnx_value(v, k)) if idx < len(self.outputs): _apply_tensor_info_to_value_info(onnx_outputs[-1], self.outputs[idx]) apply_dynamic_axes_info(onnx_outputs[-1], k) graph = onnx.helper.make_graph( nodes=onnx_nodes, name=self.traced.original_name, inputs=onnx_inputs, outputs=onnx_outputs, initializer=[v for k, v in onnx_vars.items()], doc_string=None if self.strip_doc_string else self.graph_doc_string, # TODO(twata): Use torch IR's value type info # value_info=[ # self.values[k] for k in set(list(self.values.keys())) - set(inout_names) # ], ) self.log("ONNX printable graph", onnx.helper.printable_graph(graph)) model: onnx.ModelProto = onnx.helper.make_model( graph, opset_imports=[onnx.helper.make_opsetid("", self.opset_version)], ) model = self.check_model(model) # Applying dynamic axes after onnx shape inference since it will be erased for o in model.graph.output: apply_dynamic_axes_info(o, o.name) return model