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
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_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 initializer_is_output_model():
path = os.path.join(TEST_ROOT, "models", "initializer_is_output.onnx")
model = onnx.load(path)
X = Constant(name="X", values=np.ones((64, 64), dtype=np.float32))
return Model(path,
inputs=[],
outputs=[X],
nodes=[],
opset=OnnxImporter.get_opset(model))
def identity_model():
path = os.path.join(TEST_ROOT, "models", "identity.onnx")
model = onnx.load(path)
x = Variable(name="x", dtype=np.float32, shape=(1, 1, 2, 2))
y = Variable(name="y", dtype=np.float32, shape=(1, 1, 2, 2))
node = Node(op="Identity", inputs=[x], outputs=[y])
return Model(path,
inputs=[x],
outputs=[y],
nodes=[node],
opset=OnnxImporter.get_opset(model))
def dim_param_model():
path = os.path.join(TEST_ROOT, "models", "dim_param.onnx")
model = onnx.load(path)
x = Variable(name="Input:0", dtype=np.float32, shape=("dim0", 16, 128))
y = Variable(name="Output:0", dtype=np.float32, shape=("dim0", 16, 128))
node = Node(op="Identity", inputs=[x], outputs=[y])
return Model(path,
inputs=[x],
outputs=[y],
nodes=[node],
opset=OnnxImporter.get_opset(model))
def scan_model():
path = os.path.join(TEST_ROOT, "models", "scan.onnx")
model = onnx.load(path)
# Body graph
sum_in = Variable(name="sum_in", dtype=np.float32, shape=(2, ))
next = Variable(name="next", dtype=np.float32, shape=(2, ))
sum_out = Variable(name="sum_out", dtype=np.float32, shape=(2, ))
scan_out = Variable(name="scan_out", dtype=np.float32, shape=(2, ))
body_nodes = [
Node(op="Add", inputs=[sum_in, next], outputs=[sum_out]),
Node(op="Identity", inputs=[sum_out], outputs=[scan_out]),
]
body_graph = Graph(nodes=body_nodes,
inputs=[sum_in, next],
outputs=[sum_out, scan_out],
name="scan_body")
# Outer graph
inputs = [
Variable(name="initial", dtype=np.float32, shape=(2, )),
Variable(name="x", dtype=np.float32, shape=(3, 2)),
]
outputs = [
Variable(name="y", dtype=np.float32, shape=(2, )),
Variable(name="z", dtype=np.float32, shape=(3, 2)),
]
attrs = OrderedDict()
attrs["body"] = body_graph
attrs["num_scan_inputs"] = 1
scan_node = Node(op="Scan", inputs=inputs, outputs=outputs, attrs=attrs)
return Model(
path,
inputs=inputs,
outputs=outputs,
nodes=[scan_node],
opset=OnnxImporter.get_opset(model),
)
def lstm_model():
path = os.path.join(TEST_ROOT, "models", "lstm.onnx")
model = onnx.load(path)
onnx_graph = model.graph
def load_initializer(index: int) -> np.ndarray:
return onnx.numpy_helper.to_array(onnx_graph.initializer[index])
# Optional inputs are represented by empty tensors
X = Variable(name="X", dtype=np.float32, shape=(4, 3, 6))
W = Constant(name="W", values=load_initializer(0))
R = Constant(name="R", values=load_initializer(1))
B = Constant(name="B", values=load_initializer(2))
initial_c = Constant(name="initial_c", values=load_initializer(3))
Y = Variable(name="Y", dtype=np.float32, shape=(4, 1, 3, 5))
Y_h = Variable(name="Y_h", dtype=np.float32, shape=(1, 3, 5))
Y_c = Variable(name="Y_c", dtype=np.float32, shape=(1, 3, 5))
attrs = OrderedDict()
attrs["direction"] = "forward"
attrs["hidden_size"] = 5
node = Node(
op="LSTM",
attrs=attrs,
inputs=[X, W, R, B,
Variable.empty(),
Variable.empty(), initial_c],
outputs=[Y, Y_h, Y_c],
)
# Initializers will not be included in the graph inputs.
return Model(
path,
inputs=[X],
outputs=[Y, Y_h, Y_c],
nodes=[node],
opset=OnnxImporter.get_opset(model),
)
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
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_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
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(self, model):
graph = OnnxImporter.import_graph(model.load().graph)
model.assert_equal(graph)
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
def test_import_graph_with_dim_param(self):
model = dim_param_model()
graph = OnnxImporter.import_graph(model.load().graph)
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_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
def test_export(self):
with tempfile.NamedTemporaryFile() as f:
onnx_model = gs.export_onnx(self.imported_graph)
assert onnx_model
assert OnnxImporter.import_graph(
onnx_model.graph) == self.imported_graph
def setup_method(self):
self.imported_graph = OnnxImporter.import_graph(
identity_model().load().graph)
