def test_if_with_sequence(self):
# S = [a]
# if cond is True
# S = [a,b]
# else
# S = [a,c]
a = np.random.randn(2, 1, 2).astype(np.float32)
b = np.random.randn(1, 1, 2).astype(np.float32)
c = np.random.randn(3, 1, 2).astype(np.float32)
seq_construct_node = helper.make_node('SequenceConstruct', ['a'], ['S'])
seq_insert_node1 = helper.make_node('SequenceInsert', ['S', 'b'], ['Sb'])
seq_insert_node2 = helper.make_node('SequenceInsert', ['S', 'c'], ['Sc'])
a_in = helper.make_tensor_value_info('a', onnx.TensorProto.FLOAT, [2, 1, 2])
b_in = helper.make_tensor_value_info('b', onnx.TensorProto.FLOAT, [1, 1, 2])
c_in = helper.make_tensor_value_info('c', onnx.TensorProto.FLOAT, [3, 1, 2])
cond_in = helper.make_tensor_value_info('cond', TensorProto.BOOL, [])
s_in = helper.make_sequence_value_info('S', TensorProto.FLOAT,
[None, None, None, None])
sb_out = helper.make_sequence_value_info('Sb', TensorProto.FLOAT,
[None, None, None, None])
sc_out = helper.make_sequence_value_info('Sc', TensorProto.FLOAT,
[None, None, None, None])
s_final_out = helper.make_sequence_value_info('S_final', TensorProto.FLOAT,
[None, None, None, None])
then_graph = helper.make_graph(nodes=[seq_insert_node1],
name="then_graph",
inputs=[s_in, b_in],
outputs=[sb_out])
else_graph = helper.make_graph(nodes=[seq_insert_node2],
name="else_graph",
inputs=[s_in, c_in],
outputs=[sc_out])
if_node = helper.make_node('If', ['cond'], ['S_final'],
then_branch=then_graph,
else_branch=else_graph)
graph_def = helper.make_graph(nodes=[seq_construct_node, if_node],
name='test_if',
inputs=[a_in, b_in, c_in, cond_in],
outputs=[s_final_out])
tf_rep = prepare(helper.make_model(graph_def))
output = tf_rep.run({
'a': a,
'b': b,
'c': c,
'cond': np.array(True, dtype=np.bool)
})
np.testing.assert_almost_equal(output['S_final'].values[:2], a)
np.testing.assert_almost_equal(output['S_final'].values[2:], b)
output = tf_rep.run({
'a': a,
'b': b,
'c': c,
'cond': np.array(False, dtype=np.bool)
})
np.testing.assert_almost_equal(output['S_final'].values[:2], a)
np.testing.assert_almost_equal(output['S_final'].values[2:], c)
def expect(node, inputs, outputs, name):
ginputs = [
make_sequence_value_info(node.input[0], TensorProto.FLOAT, []), # pylint: disable=E1101,
make_sequence_value_info(node.input[1], TensorProto.FLOAT, []), # pylint: disable=E1101,
]
if len(node.input) > 2:
ginputs.append(
make_tensor_value_info(node.input[2], TensorProto.INT64,
[]), # pylint: disable=E1101
)
goutputs = [
make_sequence_value_info(node.output[0], TensorProto.FLOAT,
[]), # pylint: disable=E1101,
]
model_def = make_model(
opset_imports=[make_operatorsetid('', TARGET_OPSET)],
graph=make_graph(name=name,
inputs=ginputs,
outputs=goutputs,
nodes=[node]))
oinf = OnnxInference(model_def)
got = oinf.run({n: v for n, v in zip(node.input, inputs)})
self.assertEqual(len(got), 1)
oseq = got['output_sequence']
self.assertEqual(len(oseq), len(outputs))
for e, g in zip(outputs, oseq):
self.assertEqualArray(e, g)
del model_def.opset_import[:] # pylint: disable=E1101
op_set = model_def.opset_import.add() # pylint: disable=E1101
op_set.domain = ''
op_set.version = 15
model_def.ir_version = 8
def test_loop_with_sequence(self):
if legacy_opset_pre_ver(14):
raise unittest.SkipTest(
"ONNX version {} doesn't support helper.make_tensor_sequence_value_info."
.format(defs.onnx_opset_version()))
# construct sequence S with tensor a in it
# insert tensor b into sequence S for 3 time
a = np.random.randn(2, 1, 2).astype(np.float32)
b = np.random.randn(1, 1, 2).astype(np.float32)
M = np.array(3, dtype=np.int64)
cond = np.array(True, dtype=np.bool)
seq_construct_node = helper.make_node('SequenceConstruct', ['a'],
['S'])
seq_insert_node = helper.make_node('SequenceInsert', ['S', 'b'],
['Updated_S'])
a_in = helper.make_tensor_value_info('a', onnx.TensorProto.FLOAT,
[2, 1, 2])
b_in = helper.make_tensor_value_info('b', onnx.TensorProto.FLOAT,
[1, 1, 2])
M_in = helper.make_tensor_value_info('M', TensorProto.INT64, [])
cond_init_in = helper.make_tensor_value_info('cond_init',
TensorProto.BOOL, [])
iter_count_in = helper.make_tensor_value_info('iter_count',
TensorProto.INT64, [])
cond_in = helper.make_tensor_value_info('cond', TensorProto.BOOL, [])
s_in = helper.make_sequence_value_info('S', TensorProto.FLOAT,
[None, None, None, None])
cond_out = helper.make_tensor_value_info('cond', TensorProto.BOOL, [])
s_out = helper.make_sequence_value_info('Updated_S', TensorProto.FLOAT,
[None, None, None, None])
s_final_out = helper.make_sequence_value_info('S_final',
TensorProto.FLOAT,
[None, None, None, None])
body_graph = helper.make_graph(nodes=[seq_insert_node],
name="for_loop_graph",
inputs=[iter_count_in, cond_in, s_in],
outputs=[cond_out, s_out])
loop_node = helper.make_node('Loop', ['M', '', 'S'], ['S_final'],
body=body_graph)
graph_def = helper.make_graph(nodes=[seq_construct_node, loop_node],
name='test_loop',
inputs=[a_in, b_in, M_in, cond_init_in],
outputs=[s_final_out])
tf_rep = prepare(helper.make_model(graph_def))
output = tf_rep.run({'a': a, 'b': b, 'M': M, 'cond_init': cond})
np.testing.assert_almost_equal(output['S_final'][0], a)
np.testing.assert_almost_equal(output['S_final'][1], b)
np.testing.assert_almost_equal(output['S_final'][2], b)
np.testing.assert_almost_equal(output['S_final'][3], b)
def make_onnx_inputs_outputs(name, elem_type, shape, **kwargs):
"""Wrapper for creating onnx graph inputs or outputs
name, # type: Text
elem_type, # type: TensorProto.DataType
shape, # type: Optional[Sequence[int]]
"""
if elem_type is None:
elem_type = onnx_pb.TensorProto.UNDEFINED
elif isinstance(elem_type, SeqType):
return helper.make_sequence_value_info(name, elem_type.dtype,
make_onnx_shape(shape),
**kwargs)
return helper.make_tensor_value_info(name, elem_type,
make_onnx_shape(shape), **kwargs)
def create_batch_output_sequence(model: onnx.ModelProto,
output_name: str = 'output',
output_prefix: str = 'output_'):
"""
"""
outputs = get_outputs(model)
batch_outputs = list(
filter(lambda x: x.name.startswith(output_prefix), outputs))
print(batch_outputs)
get_node_name = lambda x: x.name
batch_output_names = list(map(get_node_name, batch_outputs))
## squeeze outputs before construct sequence
## TODO : read dim first
# squeezed_output_nodes = []
# squeezed_output_names = []
# for output in batch_outputs :
# squeezed_output_name = 'squeezed_{}'.format(output.name)
# squeezed_output_nodes.append(helper.make_node(
# 'Squeeze',
# inputs=[output.name], axes=[0],
# outputs=[squeezed_output_name],
# ))
# squeezed_output_names.append(squeezed_output_name)
## create sequence info and node
sequence_value_info = helper.make_sequence_value_info(
name=output_name,
shape=None,
elem_type=onnx.TensorProto.FLOAT,
)
# sequence_construct_node = helper.make_node(
# 'SequenceConstruct',
# inputs=squeezed_output_names,
# outputs=[output_name]
# )
sequence_construct_node = helper.make_node('SequenceConstruct',
inputs=batch_output_names,
outputs=[output_name])
# for node in squeezed_output_nodes :
# model.graph.node.append(node)
model.graph.node.append(sequence_construct_node)
model.graph.output.append(sequence_value_info)
return model
def test_loop_additional_input(self):
# Given a tensor x of values [x1, ..., xN],
# Return a sequence of tensors of
# [[x1], [x1, x2], ..., [x1, ..., xN]]
cond_in = make_tensor_value_info('cond_in', TensorProto.BOOL, []) # pylint: disable=E1101
cond_out = make_tensor_value_info('cond_out', TensorProto.BOOL, []) # pylint: disable=E1101
iter_count = make_tensor_value_info('iter_count', TensorProto.INT64,
[]) # pylint: disable=E1101
seq_in = make_tensor_sequence_value_info('seq_in', TensorProto.FLOAT,
None) # pylint: disable=E1101
seq_out = make_tensor_sequence_value_info('seq_out', TensorProto.FLOAT,
None) # pylint: disable=E1101
x = numpy.array([1, 2, 3, 4, 5]).astype(numpy.float32)
x_const_node = make_node(
'Constant',
inputs=[],
outputs=['x'],
value=make_tensor(
name='const_tensor_x',
data_type=TensorProto.FLOAT, # pylint: disable=E1101
dims=x.shape,
vals=x.flatten().astype(float)))
zero_const_node = make_node(
'Constant',
inputs=[],
outputs=['slice_start'],
value=make_tensor(
name='const_tensor_zero',
data_type=TensorProto.INT64, # pylint: disable=E1101
dims=(1, ),
vals=[0]))
axes_node = make_node(
'Constant',
inputs=[],
outputs=['axes'],
value=make_tensor(
name='const_tensor_axes',
data_type=TensorProto.INT64, # pylint: disable=E1101
dims=(),
vals=[0]))
add_node = make_node('Add',
inputs=['iter_count', 'XI'],
outputs=['slice_end'])
slice_node = make_node('Slice',
inputs=['x', 'slice_start', 'slice_end'],
outputs=['slice_out'])
insert_node = make_node('SequenceInsert',
inputs=['seq_in', 'slice_out'],
outputs=['seq_out'])
identity_node = make_node('Identity',
inputs=['cond_in'],
outputs=['cond_out'])
loop_body = make_graph([
identity_node, x_const_node, zero_const_node, add_node, axes_node,
slice_node, insert_node
], 'loop_body', [iter_count, cond_in, seq_in], [cond_out, seq_out])
node = make_node('Loop',
inputs=['trip_count', 'cond', 'seq_empty'],
outputs=['seq_res'],
body=loop_body)
node1 = make_node('Neg', inputs=['XI'], outputs=['Y'])
node_concat = make_node('ConcatFromSequence',
inputs=['seq_res'],
outputs=['res'],
axis=0,
new_axis=0)
trip_count = numpy.array(5).astype(numpy.int64)
seq_empty = [] # type: List[Any]
cond = numpy.array(1).astype(numpy.bool)
model_def = make_model(
opset_imports=[make_operatorsetid('', TARGET_OPSET)],
graph=make_graph(
name='loop_test',
inputs=[
make_tensor_value_info('trip_count', TensorProto.INT64,
trip_count.shape), # pylint: disable=E1101
make_tensor_value_info('cond', TensorProto.BOOL,
cond.shape), # pylint: disable=E1101
make_sequence_value_info('seq_empty', TensorProto.FLOAT,
[]), # pylint: disable=E1101
make_tensor_value_info('XI', TensorProto.INT64, [])
], # pylint: disable=E1101
outputs=[
make_tensor_value_info('res', TensorProto.FLOAT, None), # pylint: disable=E1101
make_tensor_value_info('Y', TensorProto.INT64, [])
], # pylint: disable=E1101
nodes=[node1, node, node_concat]))
del model_def.opset_import[:] # pylint: disable=E1101
op_set = model_def.opset_import.add() # pylint: disable=E1101
op_set.domain = ''
op_set.version = 15
model_def.ir_version = 8
expected = numpy.array(
[1., 1., 2., 1., 2., 3., 1., 2., 3., 4., 1., 2., 3., 4., 5.],
dtype=numpy.float32)
X = numpy.array([1], dtype=numpy.int64)
for rt in ['python', 'onnxruntime1', 'python_compiled']:
with self.subTest(rt=rt):
oinf = OnnxInference(model_def, runtime=rt)
inputs = {
'trip_count': trip_count,
'cond': cond,
'seq_empty': seq_empty,
'XI': X
}
if rt == 'python_compiled':
code = str(oinf)
self.assertIn("context={'XI': XI}", code)
got = oinf.run(inputs)
self.assertEqualArray(-X, got['Y'])
self.assertEqualArray(expected, got['res'])
if rt == 'python':
siz = oinf.infer_sizes(inputs)
self.assertIsInstance(siz, dict)
typ = oinf.infer_types()
self.assertEqual(typ["trip_count"], numpy.int64)
if 'cond' in typ:
self.assertEqual(typ["cond"], numpy.bool_)
for k, v in typ.items():
if k in {'trip_count', 'cond', 'Y', 'XI'}:
continue
self.assertIsInstance(v, SequenceType)
def make_tensor_sequence_value_info(name, tensor_type, shape):
return make_sequence_value_info(name, tensor_type, shape, None)