def test_export_node(self):
name = "TestNode"
op = "Test"
inputs = [Variable(name="input")]
outputs = [Variable(name="output")]
attrs = OrderedDict()
attrs["float_attr"] = 4.0
attrs["int_attr"] = 10
attrs["str_attr"] = "constant"
attrs["tensor_attr"] = Constant(
"testTensor", np.ones(shape=(1, 2, 3, 4), dtype=np.float32))
attrs["floats_attr"] = [1.0, 2.0, 3.0, 4.0]
attrs["ints_attr"] = [4, 3, 2, 1]
attrs["strings_attr"] = ["constant", "and", "variable"]
node = Node(op=op,
name=name,
inputs=inputs,
outputs=outputs,
attrs=attrs)
onnx_node = OnnxExporter.export_node(node)
assert onnx_node.name == name
assert onnx_node.op_type == op
assert onnx_node.input == ["input"]
assert onnx_node.output == ["output"]
for onnx_attr, (name, attr) in zip(onnx_node.attribute, attrs.items()):
assert onnx_attr.name == name
if isinstance(attr, float):
assert onnx_attr.f == attr
elif isinstance(attr, int):
assert onnx_attr.i == attr
elif isinstance(attr, str):
assert onnx_attr.s.decode() == attr
elif isinstance(attr, Tensor):
assert onnx_attr.t.SerializeToString(
) == OnnxExporter.export_tensor_proto(
attr).SerializeToString()
elif isinstance(attr, list):
if isinstance(attr[0], float):
assert onnx_attr.floats == attr
elif isinstance(attr[0], int):
assert onnx_attr.ints == attr
elif isinstance(attr[0], str):
assert [s.decode() for s in onnx_attr.strings] == attr
else:
raise AssertionError(
"Unrecognized list attribute: ({:}: {:}) of type: {:}".
format(name, attr, type(attr)))
else:
raise AssertionError(
"Unrecognized attribute: ({:}: {:}) of type: {:}".format(
name, attr, type(attr)))
def export_onnx(graph: Graph,
do_type_check=True,
**kwargs) -> "onnx.ModelProto":
"""
Exports an onnx-graphsurgeon Graph to an ONNX model.
Args:
graph (Graph): The graph to export
Optional Args:
do_type_check (bool): Whether to check that input and output tensors have data types defined, and fail if not.
**kwargs: Additional arguments to onnx.helper.make_model
Returns:
onnx.ModelProto: A corresponding ONNX model.
"""
from onnx_graphsurgeon.exporters.onnx_exporter import OnnxExporter
import onnx
onnx_graph = OnnxExporter.export_graph(graph, do_type_check=do_type_check)
if "opset_imports" not in kwargs:
kwargs["opset_imports"] = [onnx.helper.make_opsetid("", graph.opset)]
return onnx.helper.make_model(onnx_graph, **kwargs)
def test_export_variable_tensor_empty_shape(self):
shape = None
tensor = Variable(dtype=np.float32,
shape=shape,
name="variable_tensor")
onnx_tensor = OnnxExporter.export_value_info_proto(tensor,
do_type_check=True)
assert not onnx_tensor.type.tensor_type.HasField("shape")
def test_export_graph(self, model):
onnx_graph = model.load().graph
graph = OnnxImporter.import_graph(onnx_graph)
exported_onnx_graph = OnnxExporter.export_graph(graph)
imported_graph = OnnxImporter.import_graph(exported_onnx_graph)
assert graph == imported_graph
assert graph.opset == imported_graph.opset
# ONNX exports the initializers in this model differently after importing - ONNX GS can't do much about this.
if model.path != lstm_model().path:
assert onnx_graph == exported_onnx_graph
def test_export_constant_tensor_to_tensor_proto(self):
name = "constant_tensor"
shape = (3, 224, 224)
values = np.random.random_sample(size=shape).astype(np.float32)
tensor = Constant(name=name, values=values)
onnx_tensor = OnnxExporter.export_tensor_proto(tensor)
assert onnx_tensor.name == name
assert np.all(onnx.numpy_helper.to_array(onnx_tensor) == values)
assert onnx_tensor.data_type == onnx.TensorProto.FLOAT
assert tuple(onnx_tensor.dims) == shape
def test_export_constant_tensor_lazy_values_to_tensor_proto(self):
name = "constant_tensor"
shape = (3, 3, 3)
dtype = np.float32
onnx_tensor = onnx.numpy_helper.from_array(
np.ones(shape=shape, dtype=dtype))
tensor = Constant(name=name, values=LazyValues(onnx_tensor))
# Exporter should *not* load LazyValues into a numpy array.
onnx_tensor = OnnxExporter.export_tensor_proto(tensor)
assert isinstance(tensor._values, LazyValues)
def test_export_variable_tensor_scalar_shape(self):
shape = [None]
tensor = Variable(dtype=np.float32,
shape=shape,
name="variable_tensor")
onnx_tensor = OnnxExporter.export_value_info_proto(tensor,
do_type_check=True)
assert not onnx_tensor.type.tensor_type.shape.dim[0].HasField(
"dim_param")
assert not onnx_tensor.type.tensor_type.shape.dim[0].HasField(
"dim_value")
def test_export_variable_tensor_empty_dim_param(self):
shape = ("", 224, 224)
tensor = Variable(dtype=np.float32,
shape=shape,
name="variable_tensor")
onnx_tensor = OnnxExporter.export_value_info_proto(tensor,
do_type_check=True)
onnx_shape = []
for dim in onnx_tensor.type.tensor_type.shape.dim:
onnx_shape.append(
dim.dim_value if dim.HasField("dim_value") else dim.dim_param)
assert tuple(onnx_shape) == shape
def test_export_variable_tensor(self):
name = "variable_tensor"
shape = (3, 224, 224)
dtype = np.float32
tensor = Variable(dtype=dtype, shape=shape, name=name)
onnx_tensor = OnnxExporter.export_value_info_proto(tensor,
do_type_check=True)
assert onnx_tensor.name == name
assert onnx_tensor.type.tensor_type.elem_type == onnx.TensorProto.FLOAT
onnx_shape = []
for dim in onnx_tensor.type.tensor_type.shape.dim:
onnx_shape.append(dim.dim_value)
assert tuple(onnx_shape) == shape
def test_export_constant_tensor_to_value_info_proto(self):
name = "constant_tensor"
shape = (3, 224, 224)
values = np.random.random_sample(size=shape).astype(np.float32)
tensor = Constant(name=name, values=values)
onnx_tensor = OnnxExporter.export_value_info_proto(tensor,
do_type_check=True)
assert onnx_tensor.name == name
assert onnx_tensor.type.tensor_type.elem_type == onnx.TensorProto.FLOAT
onnx_shape = []
for dim in onnx_tensor.type.tensor_type.shape.dim:
onnx_shape.append(dim.dim_value)
assert tuple(onnx_shape) == shape