def test_export_external_converters_overwrite(tmpdir, check_model_expect):
path = str(tmpdir)
model = chainer.Sequential(chainer.functions.sigmoid)
x = input_generator.positive_increasing(2, 5)
def custom_converter(params):
return onnx_helper.make_node('Tanh', params.input_names,
params.output_names),
addon_converters = {'Sigmoid': custom_converter}
export_testcase(model, x, path, external_converters=addon_converters)
tanh_outputs = chainer.functions.tanh(x).array
output_path = os.path.join(path, 'test_data_set_0', 'output_0.pb')
onnx_helper.write_tensor_pb(output_path, '', tanh_outputs) # overwrite
check_model_expect(path)
def export_testcase(model, args, out_dir, output_grad=False, **kwargs):
"""Export model and I/O tensors of the model in protobuf format.
Similar to the `export` function, this function first performs a forward
computation to a given input for obtaining an output. Then, this function
saves the pair of input and output in Protobuf format, which is a
defacto-standard format in ONNX.
This function also saves the model with the name "model.onnx".
Args:
model (~chainer.Chain): The model object.
args (list): The arguments which are given to the model
directly. Unlike `export` function, only `list` type is accepted.
out_dir (str): The directory name used for saving the input and output.
output_grad (bool): If True, this function will output model's
gradient with names 'gradient_%d.pb'.
**kwargs (dict): keyword arguments for ``onnx_chainer.export``.
"""
os.makedirs(out_dir, exist_ok=True)
model.cleargrads()
onnx_model, inputs, outputs = export(
model, args, filename=os.path.join(out_dir, 'model.onnx'),
return_named_inout=True, **kwargs)
test_data_dir = os.path.join(out_dir, 'test_data_set_0')
os.makedirs(test_data_dir, exist_ok=True)
for i, (name, var) in enumerate(inputs.items()):
pb_name = os.path.join(test_data_dir, 'input_{}.pb'.format(i))
array = chainer.cuda.to_cpu(var.array)
write_tensor_pb(pb_name, name, array)
for i, (name, var) in enumerate(outputs.items()):
pb_name = os.path.join(test_data_dir, 'output_{}.pb'.format(i))
array = chainer.cuda.to_cpu(var.array)
write_tensor_pb(pb_name, name, array)
if output_grad:
# Perform backward computation
if len(outputs) > 1:
outputs = chainer.functions.identity(*outputs)
for out in outputs.values():
out.grad = model.xp.ones_like(out.array)
list(outputs.values())[0].backward()
for i, (name, param) in enumerate(model.namedparams()):
pb_name = os.path.join(test_data_dir, 'gradient_{}.pb'.format(i))
grad = chainer.cuda.to_cpu(param.grad)
onnx_name = cleanse_param_name(name)
if grad is None:
warnings.warn(
'Parameter `{}` does not have gradient value'.format(name))
else:
write_tensor_pb(pb_name, onnx_name, grad)