class TestOnnxExporter(object):
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_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
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
# TODO: Test subgraph export.
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)))
# See test_importers for import correctness checks
# This function first imports an ONNX graph, and then re-exports it with no changes.
# The exported ONNX graph should exactly match the original.
@pytest.mark.parametrize("model", [
identity_model(),
lstm_model(),
scan_model(),
dim_param_model(),
initializer_is_output_model()
])
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_import_graph_with_dim_param(self):
model = dim_param_model()
graph = OnnxImporter.import_graph(model.load().graph)
model.assert_equal(graph)
class TestOnnxImporter(object):
def test_import_variable_tensor(self):
name = "test0"
shape = [1, 2, 3, 4]
onnx_tensor = onnx.helper.make_tensor_value_info(name, onnx.TensorProto.FLOAT, shape)
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Variable
assert tensor.name == name
assert tensor.dtype == np.float32
assert tensor.shape == shape
def test_import_constant_tensor(self):
shape = (3, 3, 3)
dtype = np.float32
onnx_tensor = onnx.numpy_helper.from_array(np.ones(shape=shape, dtype=dtype))
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Constant
assert tensor.dtype == dtype
assert tensor.shape == shape
def test_import_tensor_unknown_metadata(self):
name = "test0"
onnx_tensor = onnx.helper.make_empty_tensor_value_info(name)
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Variable
assert tensor.name == name
# TODO: Test all attribute types - missing graph
def test_import_node(self):
op = "Test"
inputs = ["x"]
outputs = ["y"]
float_attr = 4.0
int_attr = 10
str_attr = "constant"
tensor_vals = np.ones(shape=(1, 2, 3, 4), dtype=np.float32)
tensor_attr = onnx.numpy_helper.from_array(tensor_vals)
floats_attr = [1.0, 2.0, 3.0, 4.0]
ints_attr = [4, 3, 2, 1]
strings_attr = ["constant", "and", "variable"]
onnx_node = onnx.helper.make_node(op, inputs, outputs, float_attr=float_attr, int_attr=int_attr, str_attr=str_attr, tensor_attr=tensor_attr, floats_attr=floats_attr, ints_attr=ints_attr, strings_attr=strings_attr)
node = OnnxImporter.import_node(onnx_node, OrderedDict(), OrderedDict())
assert node.op == op
assert node.attrs["float_attr"] == float_attr
assert node.attrs["int_attr"] == int_attr
assert node.attrs["str_attr"] == str_attr
# Tensor should turn into a Constant
assert np.all(node.attrs["tensor_attr"].values == tensor_vals)
assert node.attrs["floats_attr"] == floats_attr
assert node.attrs["ints_attr"] == ints_attr
assert node.attrs["strings_attr"] == strings_attr
@pytest.mark.parametrize("model",
[identity_model(), lstm_model(), scan_model(), dim_param_model(), initializer_is_output_model(), nested_dup_names()],
ids=lambda model: str(model)
)
def test_import_graph(self, model):
graph = OnnxImporter.import_graph(model.load().graph)
model.assert_equal(graph)
def test_import_graph_value_info(self):
model = onnx.shape_inference.infer_shapes(identity_model().load())
graph = OnnxImporter.import_graph(model.graph)
tensors = graph.tensors()
assert all([type(tensor) == Variable and tensor.dtype is not None and tensor.shape for tensor in tensors.values()])
def test_import_graph_tensor_map_preserved(self):
model = identity_model()
tensor_map = OrderedDict()
graph = OnnxImporter.import_graph(model.load().graph, tensor_map=tensor_map)
assert len(tensor_map) == 0
model.assert_equal(graph)
def test_import_graph_with_initializer(self):
model = lstm_model()
graph = OnnxImporter.import_graph(model.load().graph)
model.assert_equal(graph)
def test_import_graph_with_dim_param(self):
model = dim_param_model()
graph = OnnxImporter.import_graph(model.load().graph)
model.assert_equal(graph)
class TestOnnxImporter(object):
def test_import_variable_tensor(self):
name = "test0"
shape = (1, 2, 3, 4)
onnx_tensor = onnx.helper.make_tensor_value_info(name, onnx.TensorProto.FLOAT, shape)
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Variable
assert tensor.name == name
assert tensor.dtype == np.float32
assert tuple(tensor.shape) == shape
def test_import_constant_tensor(self):
shape = (3, 3, 3)
dtype = np.float32
onnx_tensor = onnx.numpy_helper.from_array(np.ones(shape=shape, dtype=dtype))
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Constant
assert tensor.dtype == dtype
assert tuple(tensor.shape) == shape
def test_import_tensor_unknown_metadata(self):
name = "test0"
onnx_tensor = onnx.helper.make_empty_tensor_value_info(name)
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Variable
assert tensor.name == name
# An empty string in `dim_param` should be treated like a dynamic dimension
def test_import_empty_dim_param_tensor(self):
shape = (1, 2, "non-empty", "")
onnx_tensor = onnx.helper.make_tensor_value_info("test0", onnx.TensorProto.FLOAT, shape)
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Variable
assert tuple(tensor.shape) == shape
# Sometimes, tensor shape is not known, in which case we shouldn't import it
def test_import_unknown_shape_tensor(self):
shape = None
onnx_tensor = onnx.helper.make_tensor_value_info("test0", onnx.TensorProto.FLOAT, shape)
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Variable
assert tensor.shape is None
# Scalars can be represented in ONNX with a dim that includes neither a dim_param nor dim_value
def test_import_empty_dim_tensor(self):
shape = (None,)
onnx_tensor = onnx.helper.make_tensor_value_info("test0", onnx.TensorProto.FLOAT, shape)
onnx_tensor.type.tensor_type.shape.dim[0].ClearField("dim_value")
onnx_tensor.type.tensor_type.shape.dim[0].ClearField("dim_param")
tensor = OnnxImporter.import_tensor(onnx_tensor)
assert type(tensor) == Variable
assert tuple(tensor.shape) == shape
# TODO: Test all attribute types - missing graph
def test_import_node(self):
op = "Test"
inputs = ["x"]
outputs = ["y"]
float_attr = 4.0
int_attr = 10
str_attr = "constant"
tensor_vals = np.ones(shape=(1, 2, 3, 4), dtype=np.float32)
tensor_attr = onnx.numpy_helper.from_array(tensor_vals)
floats_attr = [1.0, 2.0, 3.0, 4.0]
ints_attr = [4, 3, 2, 1]
strings_attr = ["constant", "and", "variable"]
onnx_node = onnx.helper.make_node(
op,
inputs,
outputs,
float_attr=float_attr,
int_attr=int_attr,
str_attr=str_attr,
tensor_attr=tensor_attr,
floats_attr=floats_attr,
ints_attr=ints_attr,
strings_attr=strings_attr,
)
node = OnnxImporter.import_node(onnx_node, OrderedDict(), OrderedDict())
assert node.op == op
assert node.attrs["float_attr"] == float_attr
assert node.attrs["int_attr"] == int_attr
assert node.attrs["str_attr"] == str_attr
# Tensor should turn into a Constant
assert np.all(node.attrs["tensor_attr"].values == tensor_vals)
assert node.attrs["floats_attr"] == floats_attr
assert node.attrs["ints_attr"] == ints_attr
assert node.attrs["strings_attr"] == strings_attr
@pytest.mark.parametrize(
"model",
[
identity_model(),
lstm_model(),
scan_model(),
dim_param_model(),
initializer_is_output_model(),
nested_dup_names(),
ext_weights(),
],
ids=lambda model: str(model),
)
def test_import_graph(self, model):
graph = OnnxImporter.import_graph(model.load().graph)
model.assert_equal(graph)
def test_import_graph_value_info(self):
model = onnx.shape_inference.infer_shapes(identity_model().load())
graph = OnnxImporter.import_graph(model.graph)
tensors = graph.tensors()
assert all(
[type(tensor) == Variable and tensor.dtype is not None and tensor.shape for tensor in tensors.values()]
)
def test_import_graph_tensor_map_preserved(self):
model = identity_model()
tensor_map = OrderedDict()
graph = OnnxImporter.import_graph(model.load().graph, tensor_map=tensor_map)
assert len(tensor_map) == 0
model.assert_equal(graph)
def test_import_graph_with_initializer(self):
model = lstm_model()
graph = OnnxImporter.import_graph(model.load().graph)
model.assert_equal(graph)
def test_import_graph_with_dim_param(self):
model = dim_param_model()
graph = OnnxImporter.import_graph(model.load().graph)
model.assert_equal(graph)