def gen_imagescaler_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
test_data = np.ones([2, 3, 4, 5], dtype=np.float32)
gb.input('input', test_data)
scale = 2.0
bias = [1., 2., 3.]
expected = test_data * scale + np.array(
bias, dtype=np.float32)[None, :, None, None]
gb.output(gb.ImageScaler(
inputs=['input'], scale=scale, bias=bias, outputs=['output']),
expected)
gb.gen_test()
def fn(test_name):
gb = onnx_script.GraphBuilder(test_name)
batch_size = 2
chan = 3
wh = 5
epsilon = 1e-4
momentum = 0.95
chainer.config.train = True
input = np.random.random((batch_size, chan, wh, wh)).astype(np.float32)
bn = L.BatchNormalization(chan, decay=momentum, eps=epsilon)
# Initialize.
bn(np.random.random((batch_size, chan, wh, wh)).astype(np.float32))
scale = bn.gamma.array.copy()
bias = bn.beta.array.copy()
running_mean = bn.avg_mean.copy()
running_var = bn.avg_var.copy()
output = bn(input)
input_v = gb.input('input', input)
scale_v = gb.input('scale', scale)
bias_v = gb.input('bias', bias)
mean_in_v = gb.input('mean_in', running_mean)
var_in_v = gb.input('var_in', running_var)
output_names = ['output', 'mean_out', 'var_out']
if save_mean_var:
output_names.extend(['saved_mean', 'saved_var'])
output_v, mean_out_v, var_out_v, *saved = gb.BatchNormalization(
inputs=[input_v, scale_v, bias_v, mean_in_v, var_in_v],
epsilon=epsilon,
momentum=momentum,
outputs=output_names)
gb.output(output_v, output)
gb.output(mean_out_v, bn.avg_mean)
gb.output(var_out_v, bn.avg_var)
if saved:
mean_out_v, var_out_v = saved
mean_out = input.mean(axis=(0, 2, 3))
var_out = input.var(axis=(0, 2, 3))
np.testing.assert_allclose(output.creator.mean, mean_out)
gb.output(mean_out_v, mean_out)
gb.output(var_out_v, var_out)
gb.gen_test()
def gen_backprop_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
i = np.array(42, np.float32)
j = np.array(99, np.float32)
i_v = gb.param('i', i)
j_v = gb.param('j', j)
r_v = gb.Mul([i_v, j_v])
gb.output(r_v, i * j)
gb.gradient(i_v, j)
gb.gradient(j_v, i)
gb.gen_test()
def gen_sequence_pad_test(test_name):
# TODO(hamaji): Rewrite with GraphBuilder's input/output.
gb = onnx_script.GraphBuilder(test_name)
inputs = [np.array(a) for a in [[1, 2, 3], [4], [5, 6]]]
gb.SequenceConstruct(inputs=[], outputs=['seq0'])
for i, input in enumerate(inputs):
gb.SequenceInsert(inputs=['seq%d' % i, 'in%d' % i],
outputs=['seq%d' % (i + 1)])
index_value = 1
index_v = gb.const([index_value])
gb.SequenceAt(
inputs=['seq3', index_v],
outputs=['lookup_result'])
gb.ChainerSequencePad(
value=-42.0,
length=4,
inputs=['seq3'],
outputs=['pad3_result'])
gb.ChainerSequencePad(
value=-42.0,
inputs=['seq2'],
outputs=['pad2_result'])
gb.ChainerSequenceLengths(
inputs=['seq3'],
outputs=['seq3_lengths_seq'])
gb.ConcatFromSequence(
inputs=['seq3_lengths_seq'],
outputs=['seq3_lengths'],
axis=0,
new_axis=1)
padded = np.array([[1, 2, 3, -42], [4, -42, -42, -42], [5, 6, -42, -42]])
outputs = [
('lookup_result', np.array([4])),
('pad3_result', padded),
('pad2_result', padded[0:2, 0:3]),
('seq3_lengths', np.array([3, 1, 2])),
]
inputs = [('in%d' % i, input) for i, input in enumerate(inputs)]
inputs_vi = [_extract_value_info(a, n) for n, a in inputs]
outputs_vi = [_extract_value_info(a, n) for n, a in outputs]
graph = onnx.helper.make_graph(
nodes=gb.nodes,
name=test_name,
inputs=inputs_vi,
outputs=outputs_vi)
gen_test(graph, inputs, outputs, name=test_name)
def gen_maxpool_cover_all_test(test_name):
# A custom attribute for Chainer/ChainerX's `cover_all` parameter.
gb = onnx_script.GraphBuilder(test_name)
input = np.random.random((1, 3, 7, 7))
input_v = gb.input('input', input)
gb.output(gb.MaxPool([input_v], kernel_shape=[3, 3], strides=[2, 2],
outputs=['not_cover_all']),
F.max_pooling_2d(input, ksize=3, stride=2, cover_all=False))
gb.output(gb.MaxPool([input_v], kernel_shape=[3, 3], strides=[2, 2],
chainer_cover_all=True,
outputs=['cover_all']),
F.max_pooling_2d(input, ksize=3, stride=2, cover_all=True))
gb.gen_test()
def gen_sieve_loop(initial_sieve_val, sieve_range_tbl):
gb = onnx_script.GraphBuilder('gen_sieve')
iter = gb.input('sieve_iter', 0)
cond = gb.input('sieve_cond', True)
sieve = gb.input('sieve_in', initial_sieve_val)
n = gb.Add([iter, gb.const(2)])
is_not_me = gb.Not([gb.Equal([sieve_range_tbl, n])])
is_multiple = gb.Equal([calc_mod(gb, sieve_range_tbl, n), gb.const(0)])
is_composite = gb.And([is_not_me, is_multiple])
sieve = gb.Or([sieve, is_composite])
gb.output(gb.const(True), True)
gb.output(sieve, initial_sieve_val)
return gb.make_graph()
def gen_sequence_extend_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
input1 = aranges(3, 4)
input2 = aranges(3, 1) * 2
seq1 = [np.squeeze(i, 0) for i in np.split(input1, 3)]
seq2 = [np.squeeze(i, 0) for i in np.split(input2, 3)]
input1_v = gb.input('input1', input1)
input2_v = gb.input('input2', input2)
seq1_v = gb.SplitToSequence([input1_v], keepdims=False)
seq2_v = gb.SplitToSequence([input2_v], keepdims=False)
gb.output(gb.ChainerSequenceExtend([seq1_v, seq2_v]), Seq(seq1 + seq2))
gb.gen_test()
def gen_sequence_range_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
num_inputs = 0
for args in [(4,), (-4,), (3, 8), (5, 2),
(1, 16, 3), (1, 17, 3), (5, -2, -1), (9, 15, -1)]:
input_vs = []
for arg in args:
input_vs.append(gb.input('input_%d' % num_inputs, arg))
num_inputs += 1
output_v = gb.ChainerSequenceRange(input_vs)
len_v = gb.ChainerSequenceSize([output_v])
expected = list(range(*args))
gb.output(len_v, len(expected))
if expected:
gb.output(output_v, Seq(expected))
gb.gen_test()
def gen_generic_getitem_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
input = aranges(4, 5, 3)
reduced = np.sum(input, 0)
input_v = gb.input('input', input)
reduced_v = gb.ReduceSum([input_v], axes=[0], keepdims=False)
seq_v = gb.ChainerSequenceSeparate(inputs=[input_v])
for i in range(-2, 4):
index_v = gb.const([i])
gb.output(gb.ChainerGenericGetItem([input_v, index_v]), input[i])
gb.output(gb.ChainerGenericGetItem([reduced_v, index_v]), reduced[i])
gb.output(gb.ChainerGenericGetItem([seq_v, index_v]), input[i])
gb.gen_test()
def gen_generic_len_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
input = aranges(4, 2, 3)
input_v = gb.input('input', input)
len0_v = gb.ChainerGenericLen([input_v])
reduced_v = gb.ReduceSum([input_v], axes=[0], keepdims=False)
len1_v = gb.ChainerGenericLen([reduced_v])
seq_v = gb.ChainerSequenceSeparate(inputs=[input_v])
len_seq_v = gb.ChainerGenericLen([seq_v])
gb.output(len0_v, input.shape[0])
gb.output(len1_v, input.shape[1])
gb.output(len_seq_v, input.shape[0])
gb.gen_test()
def gen_concat_backprop_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
i = np.array([42], np.float32)
j = np.array([99], np.float32)
i_v = gb.param('i', i)
j_v = gb.param('j', j)
concat_v = gb.Concat([i_v, j_v], axis=0)
m = np.array([2, 3], np.float32)
r_v = gb.Mul([concat_v, gb.const(m)])
r = np.concatenate([i, j]) * m
gb.output(r_v, r)
gb.gradient(i_v, np.array([2], np.float32))
gb.gradient(j_v, np.array([3], np.float32))
gb.gen_test()
def gen_sequence_io_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
input = aranges(3, 2, 4)
input_seq = [[1, 2, 3, -42], [4, -42, -42, -42], [5, 6, -42, -42]]
input_v = gb.input('input', input)
input_seq_v = gb.input('input_seq', Seq(input_seq))
split_v = gb.SplitToSequence([input_v], keepdims=False)
stack_v = gb.ConcatFromSequence([input_seq_v], axis=0, new_axis=1)
gb.output(gb.Identity([input_v]), input)
gb.output(gb.Identity([input_seq_v]), Seq(input_seq))
gb.output(split_v, Seq(list(map(np.squeeze, np.split(input, len(input))))))
gb.output(stack_v, np.stack(input_seq))
gb.gen_test()
def gen_sequence_io_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
input = aranges(3, 2, 4)
input_seq = [[1, 2, 3, -42], [4, -42, -42, -42], [5, 6, -42, -42]]
input_v = gb.input('input', input)
input_seq_v = gb.input('input_seq', Seq(input_seq))
split_v = gb.ChainerSequenceSeparate([input_v])
stack_v = gb.ChainerSequenceStack([input_seq_v])
gb.output(gb.Identity([input_v]), input)
gb.output(gb.Identity([input_seq_v]), Seq(input_seq))
gb.output(split_v, Seq(list(map(np.squeeze, np.split(input, len(input))))))
gb.output(stack_v, np.stack(input_seq))
gb.gen_test()
def gen_generic_add_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
input1 = aranges(3, 4)
input2 = aranges(3, 1) * 2
seq1 = [np.squeeze(i, 0) for i in np.split(input1, 3)]
seq2 = [np.squeeze(i, 0) for i in np.split(input2, 3)]
input1_v = gb.input('input1', input1)
input2_v = gb.input('input2', input2)
seq1_v = gb.ChainerSequenceSeparate([input1_v])
seq2_v = gb.ChainerSequenceSeparate([input2_v])
gb.output(gb.ChainerGenericAdd([input1_v, input2_v]), input1 + input2)
gb.output(gb.ChainerGenericAdd([seq1_v, seq2_v]), Seq(seq1 + seq2))
gb.output(gb.ChainerGenericAdd([input1_v, seq2_v]), input1 + input2)
gb.output(gb.ChainerGenericAdd([seq1_v, input2_v]), input1 + input2)
gb.gen_test()
def gen_composites_loop(n):
gb = onnx_script.GraphBuilder('gen_sieve')
iter = gb.input('composites_iter', 0)
cond = gb.input('composites_cond', True)
# To workaround ONNX's restriction for Loop. The number of inputs
# must be greater than 2.
dummy = gb.input('composites_dummy', True)
i = gb.Add([iter, gb.const(2)])
is_greater = gb.Greater([i, n])
mod = gb.Mul([gb.Div([i, n]), n])
is_multiple = gb.Equal([i, mod])
is_composite = gb.Mul([is_greater, is_multiple])
gb.output(gb.const(True), True)
gb.output(gb.Identity([dummy]), True)
gb.output(is_composite, True)
return gb.make_graph()
def gen_generic_getslice_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
input = aranges(4, 5, 3)
reduced = np.sum(input, 0)
input_v = gb.input('input', input)
reduced_v = gb.ReduceSum([input_v], axes=[0], keepdims=False)
seq_v = gb.SplitToSequence(inputs=[input_v], keepdims=False)
def get_slice(input_v, s):
ins = [input_v]
if s.start is not None:
v = gb.const([s.start])
ins.append(v)
if s.stop is not None:
v = gb.const([s.stop])
ins.append(v)
if s.step is not None:
v = gb.const([s.step])
ins.append(v)
return gb.ChainerGenericGetSlice(ins)
def add_test(s):
expected = input[s]
gb.output(get_slice(input_v, s), expected)
gb.output(get_slice(reduced_v, s), reduced[s])
actual_v = get_slice(seq_v, s)
if len(expected):
gb.output(gb.ConcatFromSequence([actual_v], axis=0, new_axis=1),
expected)
else:
gb.output(gb.SequenceLength([actual_v]), 0)
add_test(slice(None))
for i in range(4):
add_test(slice(i, None))
for s, e in [(1, 2), (-2, 3), (0, -2), (999, 9999)]:
add_test(slice(s, e))
for s, e, t in [(1, 4, 2), (0, 100, -1), (0, 100, -2)]:
add_test(slice(s, e, t))
gb.gen_test()
def gen_sieve_loop(initial_sieve_val):
gb = onnx_script.GraphBuilder('gen_sieve')
iter = gb.input('sieve_iter', 0)
cond = gb.input('sieve_cond', True)
sieve = gb.input('sieve', initial_sieve_val)
n = gb.Add([iter, gb.const(2)])
composites_loop = gen_composites_loop(n)
_, composites_table = gb.Loop(
[gb.const(len(initial_sieve_val)),
gb.const(True),
gb.const(True)],
body=composites_loop,
outputs=['dummy', 'composites_table'])
sieve = gb.Add([sieve, composites_table])
gb.output(gb.const(True), True)
gb.output(sieve, initial_sieve_val)
return gb.make_graph()
def gen_type_coersion_test(test_name):
# Probably, ONNX expects no type coersion happens and this test is
# not valid ONNX, but we relax the restriction.
gb = onnx_script.GraphBuilder(test_name)
iv = 42
fv = 2.3
int_v = gb.const(iv)
float_v = gb.const(fv)
gb.output(gb.Add([int_v, float_v]), iv + fv)
gb.output(gb.Add([float_v, int_v]), fv + iv)
gb.output(gb.Sub([int_v, float_v]), iv - fv)
gb.output(gb.Sub([float_v, int_v]), fv - iv)
gb.output(gb.Mul([int_v, float_v]), iv * fv)
gb.output(gb.Mul([float_v, int_v]), fv * iv)
gb.output(gb.Div([int_v, float_v]), iv / fv)
gb.output(gb.Div([float_v, int_v]), fv / iv)
gb.gen_test()
def gen_prime():
max_val = 1200
gb = onnx_script.GraphBuilder('gen_prime')
num_prime = gb.input('num_prime', 100)
initial_sieve_val = np.array([False] * max_val)
sieve_loop = gen_sieve_loop(initial_sieve_val)
sieve = gb.Loop(
[gb.const(max_val),
gb.const(True),
gb.const(initial_sieve_val)],
body=sieve_loop,
outputs=['sieve'])
gb.output(sieve, np.array(get_sieve_for_test(max_val)))
gb.gen_test(validate=True)
def test_custom_op():
gb = onnx_script.GraphBuilder('pytest_custom_op')
a = np.array(13)
b = np.array(4)
c = np.array(10)
a_v = gb.input('a', a)
b_v = gb.input('b', b)
c_v = gb.input('c', c)
def custom_func(a, b, c):
return a - b, a * b - c
y, z = custom_func(a, b, c)
y_v = 'y'
z_v = 'z'
gb.ChainerDoSomething([a_v, b_v, c_v],
outputs=[y_v, z_v],
function_name='CustomFunction')
gb.output(y_v, y)
gb.output(z_v, z)
gb.gen_test()
graph = chainer_compiler_core.load('out/pytest_custom_op/model.onnx')
params = graph.params()
input_names = graph.input_names()
output_names = graph.output_names()
assert len(input_names) == 3
assert len(output_names) == 2
xcvm = graph.compile()
inputs = {}
for n, v in [('a', a), ('b', b), ('c', c)]:
inputs[n] = chainer_compiler_core.value(chainerx.array(v))
outputs = xcvm.run(inputs, custom_funcs={'CustomFunction': custom_func})
assert len(outputs) == 2
chainerx.testing.assert_allclose(9, outputs['y'].array())
chainerx.testing.assert_allclose(42, outputs['z'].array())
def gen_primes_loop(initial_primes_val, sieve, range_tbl, num_primes):
gb = onnx_script.GraphBuilder('primes_loop')
iter = gb.input('prime_iter', 0)
cond = gb.input('prime_cond', True)
primes = gb.input('prime_primes', initial_primes_val)
np = gb.input('np', 0)
is_composite = gb.Gather([sieve, iter])
then_graph = gen_primes_loop_then(initial_primes_val, primes, np)
else_graph = gen_primes_loop_else(initial_primes_val, primes, np,
iter, range_tbl)
primes, np = gb.If([is_composite],
then_branch=then_graph,
else_branch=else_graph,
outputs=['next_primes', 'next_np'])
cond = gb.Greater([to_float(gb, num_primes), to_float(gb, np)])
gb.output(cond, True)
gb.output(primes, initial_primes_val)
gb.output(np, 0)
return gb.make_graph()
def gen_prime():
max_val = 1500
# ONNX runtime crashes when this is not a multiple of 8.
max_num_primes = 104
gb = onnx_script.GraphBuilder('gen_4949_prime')
num_primes = gb.input('num_primes', max_num_primes)
# ONNX runtime does not allow using inputs of enclosing graphs.
num_primes = gb.Identity([num_primes])
sieve_range_tbl = gen_range_tbl(gb, 'sieve', max_val)
sieve_range_tbl = gb.Add([sieve_range_tbl, gb.const(2)])
prime_range_tbl = gen_range_tbl(gb, 'prime', max_num_primes)
initial_sieve_val = np.array([False] * max_val)
sieve_loop = gen_sieve_loop(initial_sieve_val, sieve_range_tbl)
sieve = gb.Loop([gb.const(max_val),
gb.const(True),
gb.const(initial_sieve_val)],
body=sieve_loop,
outputs=['sieve'])
initial_primes_val = np.array([0] * max_num_primes)
primes_loop = gen_primes_loop(initial_primes_val,
sieve,
prime_range_tbl,
num_primes)
primes, _ = gb.Loop([gb.const(max_val),
gb.const(True),
gb.const(initial_primes_val),
gb.const(0)],
body=primes_loop,
outputs=['primes', 'num_primes_out'])
primes_val = get_4949_primes_for_test(max_num_primes)
gb.output(primes, primes_val)
gb.gen_test(validate=True)
def gen_maxpool_cover_all_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
input = np.random.random((1, 3, 7, 7))
input_v = gb.input('input', input)
# Forget shape.
squeezed_v = gb.Squeeze([input_v])
dynamic_v = gb.Unsqueeze([squeezed_v], axes=[0])
gb.output(
gb.MaxPool([input_v],
kernel_shape=[3, 3],
strides=[2, 2],
outputs=['not_cover_all']),
F.max_pooling_2d(input, ksize=3, stride=2, cover_all=False))
gb.output(
gb.MaxPool([input_v],
kernel_shape=[3, 3],
strides=[2, 2],
ceil_mode=1,
outputs=['cover_all']),
F.max_pooling_2d(input, ksize=3, stride=2, cover_all=True))
gb.output(
gb.MaxPool([dynamic_v],
kernel_shape=[3, 3],
strides=[2, 2],
outputs=['not_cover_all_dynamic']),
F.max_pooling_2d(input, ksize=3, stride=2, cover_all=False))
gb.output(
gb.MaxPool([dynamic_v],
kernel_shape=[3, 3],
strides=[2, 2],
ceil_mode=1,
outputs=['cover_all_dynamic']),
F.max_pooling_2d(input, ksize=3, stride=2, cover_all=True))
gb.gen_test()
def gen_sequence_split_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
inputs = np.array([[1, 2, 3, -42], [4, -42, -42, -42], [5, 6, -42, -42]])
lengths = np.array([3, 1, 2])
inputs_v = gb.input('input', inputs)
lengths_v = gb.input('lengths', lengths)
seq_v = gb.SplitToSequence(inputs=[inputs_v],
outputs=['seq'],
keepdims=False)
lengths_seq_v = gb.SplitToSequence(inputs=[lengths_v],
outputs=['lengths_seq'],
keepdims=False)
unpadded_v = gb.ChainerSequenceUnpad(inputs=[inputs_v, lengths_seq_v],
outputs=['unpadded'])
seq_a1_v = gb.SplitToSequence(inputs=[inputs_v],
outputs=['seq_a1'],
axis=1,
keepdims=False)
for i in range(4):
index_v = gb.const([i], name='index_%d' % i)
if i < 3:
gb.output(
gb.SequenceAt(inputs=[seq_v, index_v],
outputs=['split_result_%d' % i]), inputs[i])
gb.output(
gb.SequenceAt(inputs=[unpadded_v, index_v],
outputs=['unpad_result_%d' % i]),
inputs[i][:lengths[i]])
gb.output(
gb.SequenceAt(inputs=[seq_a1_v, index_v],
outputs=['split_a1_result_%d' % i]), inputs[:, i])
gb.gen_test()
def gen_sequence_pop_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
inputs = np.array([10, 3, 4, 7, 2, 5])
inputs_v = gb.input('input', inputs)
seq_v = gb.ChainerSequenceSeparate(inputs=[inputs_v])
pop_count = 3
for i in range(pop_count):
seq_v, pop_v = gb.ChainerSequencePop(
inputs=[seq_v], outputs=['seq_%d' % i, 'pop_%d' % i])
gb.output(pop_v, inputs[-1 - i])
# This `seq_v` is used twice, so not-optimized pass will be tested.
len1_v = gb.ChainerSequenceSize(inputs=[seq_v])
seq_v, _ = gb.ChainerSequencePop(
inputs=[seq_v],
outputs=['seq_final', 'pop_final'],
)
len2_v = gb.ChainerSequenceSize(inputs=[seq_v])
gb.output(gb.Add(inputs=[len1_v, len2_v]),
(len(inputs) - pop_count) * 2 - 1)
gb.gen_test()
def gen_primes_loop_then(initial_primes_val, primes, np):
gb = onnx_script.GraphBuilder('primes_loop_then')
gb.output(gb.Identity([primes]), initial_primes_val)
gb.output(gb.Identity([np]), 0)
return gb.make_graph()
def gen_const_int_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
c = np.array(list(range(20)))
c_v = gb.const(c)
gb.output(gb.Identity([c_v]), c)
gb.gen_test()
def gen_const_prop_use_twice_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
c = np.array(list(range(20)))
c_v = gb.const(c)
gb.output(gb.Add([c_v, c_v]), c * 2)
gb.gen_test()
def fn(test_name):
gb = onnx_script.GraphBuilder(test_name)
i = np.array([42, -24], dtype=dtype)
i_v = gb.input('input', i)
gb.output(gb.Abs([i_v]), np.abs(i))
gb.gen_test()
def gen_sequence_constants_test(test_name):
gb = onnx_script.GraphBuilder(test_name)
inputs = [4, 2, 3]
seq_v = gb.const_seq(inputs)
gb.output(seq_v, Seq(inputs))
gb.gen_test()