def dump_test_inputs_outputs(inputs, outputs, test_data_dir):
if not os.path.exists(test_data_dir):
os.makedirs(test_data_dir)
for typ, values in [('input', inputs), ('output', outputs)]:
for i, (value_info, value) in enumerate(values):
name = onnx_name(value_info)
if isinstance(value, list):
assert value
digits = len(str(len(value)))
for j, v in enumerate(value):
filename = os.path.join(
test_data_dir,
'%s_%d_%s.pb' % (typ, i, str(j).zfill(digits)))
tensor = numpy_helper.from_array(v, name)
with open(filename, 'wb') as f:
f.write(tensor.SerializeToString())
#value_info.type.CopyFrom(onnx.TypeProto())
#sequence_type = value_info.type.sequence_type
#tensor_type = sequence_type.elem_type.tensor_type
#tensor_type.elem_type = tensor.data_type
else:
filename = os.path.join(test_data_dir, '%s_%d.pb' % (typ, i))
if value is None:
if get_test_args().allow_unused_params:
continue
raise RuntimeError('Unused parameter: %s' % name)
tensor = numpy_helper.from_array(value, name)
with open(filename, 'wb') as f:
f.write(tensor.SerializeToString())
vi = onnx.helper.make_tensor_value_info(
name, tensor.data_type, tensor.dims)
def generate_testcase(model,
xs,
subname=None,
output_dir=None,
backprop=False):
if output_dir is None:
args = get_test_args()
output_dir = args.output
if backprop:
output_dir = output_dir + '_backprop'
if not _seen_subnames:
# Remove all related directories to renamed tests.
for d in [output_dir] + glob.glob(output_dir + '_*'):
if os.path.isdir(d):
shutil.rmtree(d)
assert (backprop, subname) not in _seen_subnames
_seen_subnames.add((backprop, subname))
if subname is not None:
output_dir = output_dir + '_' + subname
else:
assert subname is None
if not os.path.exists(output_dir):
os.makedirs(output_dir)
def get_model():
if isinstance(model, type) or isinstance(model, types.FunctionType):
return model()
return model
model = get_model()
chainer.config.train = backprop
model.cleargrads()
ys = model(*xs)
chainer_out = validate_chainer_output(ys)
onnxmod = compile_model(model, xs)
input_tensors = onnxmod.inputs
output_tensors = onnxmod.outputs
if backprop:
ys.grad = np.ones(ys.shape, ys.dtype)
ys.backward()
if len(output_tensors) < len(chainer_out):
assert len(output_tensors) == 1
chainer_out = [np.array(chainer_out)]
assert len(output_tensors) == len(chainer_out)
outputs = list(zip(output_tensors, chainer_out))
'''
if backprop:
for name, param in sorted(model.namedparams()):
bp_name = onnx.helper.make_tensor_value_info(
'grad_out@' + name, onnx.TensorProto.FLOAT, ())
outputs.append((bp_name, param.grad))
'''
xs = list(map(lambda x: _validate_inout(x), xs))
dump_test_inputs_outputs(list(zip(input_tensors, xs)), outputs,
os.path.join(output_dir, 'test_data_set_0'))
with open(os.path.join(output_dir, 'model.onnx'), 'wb') as fp:
fp.write(onnxmod.model.SerializeToString())
def reset_test_generator(args):
_seen_subnames.clear()
get_test_args(args)
