def rewrite_onnx_testdir(model_testdir, out_testdir, new_input_types):
os.makedirs(out_testdir, exist_ok=True)
xmodel = rewrite_onnx_file(os.path.join(model_testdir, 'model.onnx'),
os.path.join(out_testdir, 'model.onnx'),
new_input_types)
name_to_type = {}
for vi in (list(xmodel.graph.input) + list(xmodel.graph.value_info) +
list(xmodel.graph.output)):
dtype = mapping.TENSOR_TYPE_TO_NP_TYPE[vi.type.tensor_type.elem_type]
shape = [d.dim_value for d in vi.type.tensor_type.shape.dim]
name_to_type[vi.name] = Type(dtype=dtype, shape=shape)
for test_set in glob.glob(os.path.join(model_testdir, 'test_data_set_*')):
dest_dir = os.path.join(out_testdir, os.path.basename(test_set))
os.makedirs(dest_dir, exist_ok=True)
for tensor_proto in glob.glob(os.path.join(test_set, '*.pb')):
xtensor = onnx.load_tensor(tensor_proto)
if xtensor.name not in name_to_type:
raise RuntimeError('Unknown tensor name: %s' % xtensor.name)
rewrite_onnx_tensor(xtensor, name_to_type[xtensor.name])
out_tensor_proto = os.path.join(dest_dir,
os.path.basename(tensor_proto))
onnx.save_tensor(xtensor, out_tensor_proto)
def test_save_and_load_tensor(self): # type: () -> None
proto = self._simple_tensor()
cls = TensorProto
proto_string = onnx._serialize(proto)
# Test if input is string
loaded_proto = onnx.load_tensor_from_string(proto_string)
self.assertTrue(proto == loaded_proto)
# Test if input has a read function
f = io.BytesIO()
onnx.save_tensor(loaded_proto, f)
f = io.BytesIO(f.getvalue())
loaded_proto = onnx.load_tensor(f, cls)
self.assertTrue(proto == loaded_proto)
# Test if input is a file name
try:
tfile = tempfile.NamedTemporaryFile(delete=False)
onnx.save_tensor(proto, tfile)
tfile.close()
loaded_proto = onnx.load_tensor(tfile.name, cls)
self.assertTrue(proto == loaded_proto)
finally:
os.remove(tfile.name)
def test_save_and_load_tensor(self): # type: () -> None
proto = self._simple_tensor()
cls = TensorProto
proto_string = onnx._serialize(proto)
# Test if input is string
loaded_proto = onnx.load_tensor_from_string(proto_string)
self.assertTrue(proto == loaded_proto)
# Test if input has a read function
f = io.BytesIO()
onnx.save_tensor(loaded_proto, f)
f = io.BytesIO(f.getvalue())
loaded_proto = onnx.load_tensor(f, cls)
self.assertTrue(proto == loaded_proto)
# Test if input is a file name
try:
tfile = tempfile.NamedTemporaryFile(delete=False)
onnx.save_tensor(proto, tfile)
tfile.close()
loaded_proto = onnx.load_tensor(tfile.name, cls)
self.assertTrue(proto == loaded_proto)
finally:
os.remove(tfile.name)
def SaveTensorProto(file_path, variable, data, name):
# ONNX input shape always has sequence axis as the first dimension, if sequence axis exists
if len(variable.dynamic_axes) == 2:
data = data.transpose((
1,
0,
) + tuple(range(2, len(data.shape))))
tp = numpy_helper.from_array(data, name if name else variable.uid)
onnx.save_tensor(tp, file_path)
def update_name_in_pb(filename, name, output_filename):
"""Update the name of the tensor in the pb file."""
tensor = onnx.load_tensor(filename)
tensor.name = name
if not output_filename:
output_filename = filename
onnx.save_tensor(tensor, output_filename)
def save_data(test_data_dir, var_and_data, output_uid_to_ort_pairs=None):
for i, (var, data) in enumerate(var_and_data.items()):
data = np.asarray(data).astype(var.dtype)
# ONNX input shape always has sequence axis as the first dimension, if sequence axis exists
if len(var.dynamic_axes) == 2:
data = data.transpose((
1,
0,
) + tuple(range(2, len(data.shape))))
file_path = os.path.join(
test_data_dir,
'{}_{}.pb'.format('output' if output_uid_to_ort_pairs else 'input',
i))
if output_uid_to_ort_pairs:
tensor_name = output_uid_to_ort_pairs[var.uid]
else:
tensor_name = var.name if var.name else var.uid
onnx.save_tensor(numpy_helper.from_array(data, tensor_name), file_path)
def main():
from_dir = sys.argv[1]
to_dir = sys.argv[2]
os.makedirs(to_dir, exist_ok=True)
xmodel = onnx.load(os.path.join(from_dir, 'model.onnx'))
convert_model(xmodel)
onnx.save(xmodel, os.path.join(to_dir, 'model.onnx'))
for test_dir in glob.glob(os.path.join(from_dir, 'test_data_set_*')):
dir_name = os.path.basename(test_dir)
to_test_dir = os.path.join(to_dir, dir_name)
os.makedirs(to_test_dir, exist_ok=True)
for pb_filename in glob.glob(os.path.join(test_dir, '*.pb')):
pb_name = os.path.basename(pb_filename)
tensor = onnx.load_tensor(pb_filename)
convert_tensor(tensor)
onnx.save_tensor(tensor, os.path.join(to_test_dir, pb_name))
def main():
from_dir = sys.argv[1]
to_dir = sys.argv[2]
os.makedirs(to_dir, exist_ok=True)
xmodel = onnx.load(os.path.join(from_dir, 'model.onnx'))
convert_model(xmodel)
onnx.save(xmodel, os.path.join(to_dir, 'model.onnx'))
for test_dir in glob.glob(os.path.join(from_dir, 'test_data_set_*')):
dir_name = os.path.basename(test_dir)
to_test_dir = os.path.join(to_dir, dir_name)
os.makedirs(to_test_dir, exist_ok=True)
for pb_filename in glob.glob(os.path.join(test_dir, '*.pb')):
pb_name = os.path.basename(pb_filename)
tensor = onnx.load_tensor(pb_filename)
convert_tensor(tensor)
onnx.save_tensor(tensor, os.path.join(to_test_dir, pb_name))
def save_ort_test_case(ort_test_case_dir):
if not os.path.exists(ort_test_case_dir):
os.makedirs(ort_test_case_dir)
test_data_set_dir = os.path.join(ort_test_case_dir,
'test_data_set_0')
if not os.path.exists(test_data_set_dir):
os.makedirs(test_data_set_dir)
onnx.save(model, os.path.join(ort_test_case_dir, 'model.onnx'))
for i, (input_name, input) in enumerate(feeds.items()):
onnx.save_tensor(
onnx.numpy_helper.from_array(input, input_name),
os.path.join(test_data_set_dir, 'input_{0}.pb'.format(i)))
output_names = [output.name for output in model.graph.output]
output_dict = dict(zip(output_names, outputs))
for i, (output_name, output) in enumerate(output_dict.items()):
onnx.save_tensor(
onnx.numpy_helper.from_array(output, output_name),
os.path.join(test_data_set_dir, 'output_{0}.pb'.format(i)))
#!/usr/bin/env python3
import onnx
import onnx.helper
import numpy as np
values = ['a','b','c','d','cat','tiger',]
output_tensor = onnx.helper.make_tensor("values", onnx.TensorProto.STRING, [len(values)], [bytes(i, 'utf-8') for i in values])
outputs = [onnx.helper.make_tensor_value_info("values", onnx.TensorProto.STRING, [len(values)])]
node = onnx.helper.make_node(
'Constant',
inputs=[],
outputs=['values'],
value_strings=values
)
graph = onnx.helper.make_graph(
[node],
'test-constant',
[],
outputs)
model = onnx.helper.make_model(graph, producer_name='onnx-example')
onnx.save_model(model, 'model.onnx')
onnx.save_tensor(output_tensor, 'test_data_set_0/output_0.pb')
def numpy_to_pb(name, np_data, out_filename):
"""Convert numpy data to a protobuf file."""
tensor = numpy_helper.from_array(np_data, name)
onnx.save_tensor(tensor, out_filename)
assert len(inputs) == 1
input = inputs[0]
old_shape = []
new_shape = []
for i in range(len(input.type.tensor_type.shape.dim)):
nd = input.type.tensor_type.shape.dim[i].dim_value
old_shape.append(nd)
if i >= 2:
nd = int(nd * ratio)
new_shape.append(nd)
input.type.tensor_type.shape.dim[i].dim_value = nd
print('new_shape', new_shape)
onnx.save(model, os.path.join(output_dir, 'model.onnx'))
input_pb = os.path.join(output_dir, 'test_data_set_0/input_0.pb')
input_tensor = onnx.load_tensor(input_pb)
assert len(input_tensor.dims) == len(new_shape)
input = onnx.numpy_helper.to_array(input_tensor)
pad_width = []
for od, nd in zip(old_shape, new_shape):
pad_width.append((0, nd - od))
input = np.pad(input, pad_width, 'constant')
onnx.save_tensor(onnx.numpy_helper.from_array(input), input_pb)
assert len(inputs) == 1
input = inputs[0]
old_shape = []
new_shape = []
for i in range(len(input.type.tensor_type.shape.dim)):
nd = input.type.tensor_type.shape.dim[i].dim_value
old_shape.append(nd)
if i >= 2:
nd = int(nd * ratio)
new_shape.append(nd)
input.type.tensor_type.shape.dim[i].dim_value = nd
print('new_shape', new_shape)
onnx.save(model, os.path.join(output_dir, 'model.onnx'))
input_pb = os.path.join(output_dir, 'test_data_set_0/input_0.pb')
input_tensor = onnx.load_tensor(input_pb)
assert len(input_tensor.dims) == len(new_shape)
input = onnx.numpy_helper.to_array(input_tensor)
pad_width = []
for od, nd in zip(old_shape, new_shape):
pad_width.append((0, nd - od))
input = np.pad(input, pad_width, 'constant')
onnx.save_tensor(onnx.numpy_helper.from_array(input), input_pb)
def numpy_to_pb(name, np_data, out_filename):
tensor = numpy_helper.from_array(np_data, name)
onnx.save_tensor(tensor, out_filename)
print(len(tensor.raw_data))
print(dir(tensor))
#tensor.dims[:] = [672, 672]
np_array = numpy_helper.to_array(tensor)
print("Name:", tensor.name)
print("Data Type:", tensor.data_type)
print("Shape:", np_array.shape)
print(np_array)
def numpy_to_pb(name, np_data, out_filename):
tensor = numpy_helper.from_array(np_data, name)
onnx.save_tensor(tensor, out_filename)
if __name__ == '__main__':
numpy.set_printoptions(threshold=sys.maxsize)
#print_pb_file(
# "test/super_resolution/test_data_set_0/input_0.pb")
#print_pb_file(
# "test/super_resolution/test_data_set_0/output_0.pb")
tensor = onnx.load_tensor(
"test/super_resolution/test_data_set_0/output_0.pb")
tensor.dims[:] = [1, 1, 672, 672]
onnx.save_tensor(
tensor, "test/super_resolution/test_data_set_0/output_0_fixed.pb")
#print_pb_file(
# "../test/node/test_maxpool_2d_same_upper/test_data_set_0/output_0.pb")
