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
