def test_optimizer(self):
val = np.asarray([[[[1.0, 2.0, 3.0], [1.1, 2.1, 3.1]]]], np.float32)
nodes = []
nodes[0:] =\
[helper.make_node('Constant', [], ['const1'], value=helper.make_tensor(
name='const0',
data_type=onnx_proto.TensorProto.FLOAT,
dims=val.shape,
vals=val.flatten().astype(float)))]
nodes[1:] = [helper.make_node('Identity', ['const1'], ['identity1'])]
nodes[2:] = [helper.make_node('Identity', ['identity1'], ['identity2'])]
nodes[3:] = [helper.make_node('Max', ['input1', 'identity2'], ['max0'])]
nodes[4:] = [helper.make_node('Transpose', ['max0'], ['tranpose0'], perm=[0, 2, 3, 1])]
nodes[5:] = [helper.make_node('Transpose', ['tranpose0'], ['tranpose1'], perm=(0, 3, 1, 2))]
nodes[6:] = [helper.make_node('Relu', ['tranpose1'], ['output0'], perm=(0, 3, 1, 2))]
input0 = helper.make_tensor_value_info('input1', onnx_proto.TensorProto.FLOAT, [1, 1, 2, 3])
output0 = helper.make_tensor_value_info('output0', onnx_proto.TensorProto.FLOAT, [1, 1, 2, 3])
graph = helper.make_graph(nodes, 'test0', [input0], [output0])
model = helper.make_model(graph)
self.assertIsNotNone(model)
new_nodes = optimize_onnx(nodes, inputs=[input0], outputs=[output0])
self.assertEqual(len(new_nodes), 3)
graph = helper.make_graph(new_nodes, 'test0', [input0], [output0])
model = helper.make_model(graph)
self.assertIsNotNone(model)
def build_from_onnx(onnx_nodes, nchw_inputs, inputs, outputs):
view = []
var_map = {}
for o_ in onnx_nodes:
ln = LinkedNode(o_)
view.append(ln)
for var_ in o_.output:
assert var_map.get(var_) is None
var_map[var_] = ln
additional_nodes = []
for n_ in view:
for var_ in n_.origin.input:
target = var_map.get(var_)
if target is None:
assert var_ == '' or var_ in inputs
target = LinkedNode(out_n=[var_]) # create an empty node as input
new_output = var_ + '_nhwc'
if var_ in nchw_inputs:
nnode = LinkedNode(
helper.make_node(
'Transpose',
[var_],
[new_output],
perm=[0, 2, 3, 1]))
var_map[new_output] = nnode
nnode.add_precedence(target, var_)
n_.in_redirect(var_, new_output)
target = nnode
var_ = new_output
additional_nodes.append(nnode)
n_.add_precedence(target, var_)
for n_ in view: # add a dummy output node.
for var_ in n_.origin.output:
if var_ in outputs:
LinkedNode(in_n=[var_]).add_precedence(n_, var_)
return view + additional_nodes