def test_fuse_add_bias_into_conv_use_weight_shape(self): # type: () -> None
nodes = [helper.make_node("Conv", ["X", "Y"], ["Z"]),
helper.make_node("Add", ["Z", "A"], ["B"])]
nodes.extend(self._make_fake_loop_op(
[helper.make_node("Conv", ["_X", "_Y"], ["_Z"]),
helper.make_node("Add", ["_Z", "_A"], ["_B2"])],
[(TensorProto.FLOAT, (1, 5, 3, 3), "X"),
(TensorProto.FLOAT, (16, 5, 3, 3), "Y"),
(TensorProto.FLOAT, (16, 1, 1), "A")],
[(TensorProto.FLOAT, (1, 16, 3, 3), "B2")]))
graph = helper.make_graph(
nodes,
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3)),
helper.make_tensor_value_info("A", TensorProto.FLOAT, (16, 1, 1))],
[helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 3, 3))],
)
optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])
# Squeeze, Conv, Constant (trip count), Constant (condition), Loop
assert len(list(optimized_model.graph.node)) == 5
assert optimized_model.graph.node[0].op_type == 'Squeeze'
assert optimized_model.graph.node[1].op_type == 'Conv'
assert optimized_model.graph.output[0].name == 'Z'
# Squeeze, Conv
assert len(optimized_model.graph.node[4].attribute[0].g.node) == 2
assert optimized_model.graph.node[4].attribute[0].g.node[0].op_type == 'Squeeze'
assert optimized_model.graph.node[4].attribute[0].g.node[1].op_type == 'Conv'
# Output 1 since 0 is 'cond'
assert optimized_model.graph.node[4].attribute[0].g.output[1].name == '_Z'
def test_reshape_like():
in_shape = (4, 3, 3, 4)
ref_shape = (3, 4, 4, 3)
ref_array = np.random.uniform(size=ref_shape).astype('float32')
ref_node = onnx.helper.make_node('Constant',
inputs=[],
outputs=['ref_in'],
value=onnx.helper.make_tensor(name = 'const_tensor',
data_type = onnx.TensorProto.FLOAT,
dims = ref_array.shape,
vals = ref_array.flatten().astype(float)))
copy_node = helper.make_node("Identity", ["ref_in"], ["copy_in"])
reshape_node = helper.make_node("Reshape", ["in", "copy_in"], ["out"])
graph = helper.make_graph([ref_node, copy_node, reshape_node],
"reshape_like_test",
inputs = [helper.make_tensor_value_info("in",
TensorProto.FLOAT, list(in_shape))],
outputs = [helper.make_tensor_value_info("out",
TensorProto.FLOAT, list(ref_shape))])
model = helper.make_model(graph, producer_name='reshape_like_test')
for target, ctx in ctx_list():
x = np.random.uniform(size=in_shape).astype('float32')
tvm_out = get_tvm_output(model, x, target, ctx, ref_shape, 'float32')
tvm.testing.assert_allclose(ref_shape, tvm_out.shape)
def _test_upsample_bilinear_opset9():
scale = 2
in_shape = (1, 1, 3, 3)
out_shape = (1, 1, 3*scale, 3*scale)
y = helper.make_node("Upsample", ['in','scales'], ['out'], mode='linear')
scales=[1.0, 1.0, 2.0, 2.0]
in_array = np.random.uniform(size=in_shape).astype(np.float32)
out_array = topi.testing.bilinear_resize_python(in_array, (3*scale, 3*scale), "NCHW")
ref_array = np.array(scales)
ref_node = helper.make_node('Constant',
inputs=[],
outputs=['scales'],
value=onnx.helper.make_tensor(name = 'const_tensor',
data_type = TensorProto.FLOAT,
dims = ref_array.shape,
vals = ref_array.flatten().astype(float)))
graph = helper.make_graph([ref_node, y],
'upsample_bilinear_opset9_test',
inputs = [helper.make_tensor_value_info("in", TensorProto.FLOAT, list(in_shape))],
outputs = [helper.make_tensor_value_info("out", TensorProto.FLOAT, list(out_shape))])
model = helper.make_model(graph, producer_name='upsample_bilinear_opset9_test')
inputs = []
inputs.append(in_array)
for target, ctx in ctx_list():
tvm_out = get_tvm_output(model, inputs, target, ctx, out_shape, 'float32')
tvm.testing.assert_allclose(out_array, tvm_out, rtol=1e-5, atol=1e-5)
def test_nested_graph(self): # type: () -> None
n1 = helper.make_node(
"Scale", ["X"], ["Y"], scale=2., name="n1")
n2 = helper.make_node(
"Scale", ["Y"], ["Z"], scale=3., name="n2")
graph = helper.make_graph(
[n1, n2],
"nested",
inputs=[
helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2])
],
outputs=[
helper.make_tensor_value_info("Z", TensorProto.FLOAT, [1, 2])
]
)
i1 = helper.make_node(
"If", ["cond"], ["Z"], then_branch=graph, else_branch=graph)
graph = helper.make_graph(
[i1],
"test",
inputs=[
helper.make_tensor_value_info("cond", TensorProto.BOOL, [1])
],
outputs=[],
)
checker.check_graph(graph)
def verify_reduce_x(name, indata, axis, keepdims):
indata = np.array(indata).astype(np.float32)
# numpy expect result
if name == 'ReduceMax':
outdata = np.maximum.reduce(indata, axis=axis, keepdims=keepdims == 1)
elif name == 'ReduceMin':
outdata = np.minimum.reduce(indata, axis=axis, keepdims=keepdims == 1)
elif name == 'ReduceSum':
outdata = np.sum(indata, axis=axis, keepdims=keepdims == 1)
elif name == 'ReduceMean':
outdata = np.mean(indata, axis=axis, keepdims=keepdims == 1)
else:
raise Exception('unsupport op: {}'.format(name))
if len(np.asarray(outdata).shape) == 0:
outdata = np.asarray([outdata])
# onnx graph
if axis is None:
node = helper.make_node(name, inputs=['input'], outputs=['output'],
keepdims=keepdims)
else:
node = helper.make_node(name, inputs=['input'], outputs=['output'],
axis=axis, keepdims=keepdims)
graph = helper.make_graph([node],
'{}_test'.format(name),
inputs = [helper.make_tensor_value_info("input",
TensorProto.FLOAT, list(indata.shape))],
outputs = [helper.make_tensor_value_info("output",
TensorProto.FLOAT, list(outdata.shape))])
model = helper.make_model(graph, producer_name='{}_test'.format(name))
# tvm result
for target, ctx in ctx_list():
tvm_out = get_tvm_output(model, indata, target, ctx, outdata.shape, 'float32')
tvm.testing.assert_allclose(outdata, tvm_out, rtol=1e-5, atol=1e-5)
def verify_constantfill(is_shape, input_dim, out_dim, value, dtype, **kwargs):
input_a = np.random.uniform(size=input_dim).astype(dtype)
out = np.empty(shape=out_dim, dtype=dtype)
out.fill(value)
if is_shape == True:
fill_node = helper.make_node("ConstantFill", [], ["out"], shape=input_dim, value=value, **kwargs)
else:
fill_node = helper.make_node("ConstantFill", ["input_a"], ["out"], value=value, dtype=dtype, **kwargs)
graph = helper.make_graph([fill_node],
"fill_test",
inputs = [helper.make_tensor_value_info("input_a",
TensorProto.FLOAT, list(input_dim))],
outputs = [helper.make_tensor_value_info("out",
TensorProto.FLOAT, list(out.shape))])
model = helper.make_model(graph, producer_name='fill_test')
for target, ctx in ctx_list():
if is_shape == True:
tvm_out = get_tvm_output(model, [], target, ctx, out.shape)
else:
tvm_out = get_tvm_output(model, [input_a], target, ctx, out.shape)
tvm.testing.assert_allclose(out, tvm_out, rtol=1e-5, atol=1e-5)
def _make_fake_loop_op(self,
body_nodes, # type: Sequence[NodeProto]
input_types, # type: Sequence[Tuple[TensorProto.DataType, Sequence[int], Text]]
output_types # type: Sequence[Tuple[TensorProto.DataType, Sequence[int], Text]]
): # type: (...) -> List[NodeProto]
zero = helper.make_tensor("trip_count_value", TensorProto.INT32, (), [10])
true = helper.make_tensor("condition", TensorProto.BOOL, (), [True])
# lcd is a dummy loop-carried dependency that only exists because
# right now the schema checker is broken and assumes a variadic
# input needs at least one value.
graph_inputs = [helper.make_tensor_value_info("i", TensorProto.INT32, ()),
helper.make_tensor_value_info("cond", TensorProto.BOOL, ())]
for type, shape, name in input_types:
graph_inputs.append(helper.make_tensor_value_info("_" + name, type, shape))
graph_outputs = [helper.make_tensor_value_info("cond", TensorProto.BOOL, ())]
for type, shape, name in output_types:
graph_outputs.append(helper.make_tensor_value_info("_" + name, type, shape))
body_graph = helper.make_graph(body_nodes, "body_graph", graph_inputs,
graph_outputs)
loop_inputs = ["trip_count", "condition"]
loop_inputs.extend([name for _, _, name in input_types])
# TODO: fix checker to accept 0-input variadic inputs
if len(loop_inputs) == 2:
loop_inputs.append("")
loop_outputs = [name for _, _, name in output_types]
retval_nodes = [
helper.make_node("Constant", [], ["trip_count"], value=zero),
helper.make_node("Constant", [], ["condition"], value=true),
helper.make_node("Loop", loop_inputs, loop_outputs, body=body_graph)
]
return retval_nodes
def test_eliminate_unused_initializer_no_eliminate_used(self): # type: () -> None
nodes = [helper.make_node("Add", ["X", "A"], ["Z"])]
nodes.extend(self._make_fake_loop_op(
[helper.make_node("Add", ["_X", "_A"], ["_Z2"])],
[(TensorProto.FLOAT, (1, 2), "X"),
(TensorProto.FLOAT, (1, 2), "A")],
[(TensorProto.FLOAT, (1, 2), "Z2")]))
graph = helper.make_graph(
nodes,
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 2)),
helper.make_tensor_value_info("A", TensorProto.FLOAT, (1, 2))],
[helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 2))],
[helper.make_tensor("A", TensorProto.FLOAT,
dims=(1, 2),
vals=np.random.randn(1, 2).astype(np.float32).tobytes(),
raw=True)])
optimized_model = self._optimized(graph, ["eliminate_unused_initializer"])
# Add, Constant (trip count), Constant (cond), Loop
assert len(list(optimized_model.graph.node)) == 4
assert optimized_model.graph.node[0].op_type == "Add"
assert optimized_model.graph.output[0].name == "Z"
# Add
assert len(optimized_model.graph.node[3].attribute[0].g.node) == 1
assert optimized_model.graph.node[3].attribute[0].g.node[0].op_type == 'Add'
assert optimized_model.graph.node[3].attribute[0].g.output[1].name == '_Z2'
assert len(list(optimized_model.graph.initializer)) == 1
def test_fuse_add_bias_into_conv_use_conv_shape(self): # type: () -> None
sub = helper.make_node("Sub", ["M", "N"], ["Y"])
conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
add = helper.make_node("Add", ["Z", "A"], ["B"])
graph = helper.make_graph(
[sub, conv, add],
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
helper.make_tensor_value_info("M", TensorProto.FLOAT, (16, 5, 3, 3)),
helper.make_tensor_value_info("N", TensorProto.FLOAT, (16, 5, 3, 3)),
helper.make_tensor_value_info("A", TensorProto.FLOAT, (1, 16, 1, 1))],
[helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 3, 3))],
value_info=[
helper.make_tensor_value_info("Z", TensorProto.FLOAT, (1, 16, 3, 3))
],
)
optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])
assert len(optimized_model.graph.node) == 3
assert optimized_model.graph.node[0].op_type == 'Sub'
assert optimized_model.graph.node[1].op_type == 'Squeeze'
assert optimized_model.graph.node[2].op_type == 'Conv'
assert optimized_model.graph.output[0].name == 'Z'
assert optimized_model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT
assert len(optimized_model.graph.output[0].type.tensor_type.shape.dim) == 4
def test_extract_constant_to_initializer(self): # type: () -> None
conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
constant = helper.make_node("Constant", [], ["A"],
value=helper.make_tensor(
name="bias",
data_type=TensorProto.FLOAT,
dims=(16,),
vals=np.random.randn(16).astype(np.float32).tolist()))
add = helper.make_node("Add", ["Z", "A"], ["B"])
graph = helper.make_graph(
[conv, constant, add],
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3))],
[helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 3, 3))],
)
optimized_model = self._optimized(graph, ["extract_constant_to_initializer"])
self.assertEqual(
set(vi.name for vi in optimized_model.graph.input),
{'X', 'Y', 'A'})
self.assertEqual(len(optimized_model.graph.initializer), 1)
init = optimized_model.graph.initializer[0]
self.assertEqual(init.name, 'A')
self.assertEqual(init.dims, [16])
self.assertEqual(init.data_type, TensorProto.FLOAT)
self.assertEqual([n.op_type for n in optimized_model.graph.node], ['Conv', 'Add'])
def test_shape():
in_shape = (4, 3, 3, 4)
ref_shape = (6, 2, 4, 3)
ref_array = np.array(ref_shape)
ref_node = onnx.helper.make_node('Constant',
inputs=[],
outputs=['ref_in'],
value=onnx.helper.make_tensor(name = 'const_tensor',
data_type = onnx.TensorProto.INT32,
dims = ref_array.shape,
vals = ref_array.flatten().astype(int)))
reshape_node = helper.make_node("Reshape", ["in", "ref_in"], ["out"])
shape_node = helper.make_node("Shape", ['out'], ['final_out'])
graph = helper.make_graph([ref_node, reshape_node, shape_node],
"shape_test",
inputs = [helper.make_tensor_value_info("in",
TensorProto.FLOAT, list(in_shape))],
outputs = [helper.make_tensor_value_info("final_out",
TensorProto.FLOAT, list(ref_shape))])
model = helper.make_model(graph, producer_name='shape_test')
for target, ctx in ctx_list():
x = np.random.uniform(size=in_shape).astype('int32')
tvm_out = get_tvm_output(model, x, target, ctx, ref_shape, 'int32')
tvm.testing.assert_allclose(ref_shape, tvm_out)
def test_fuse_add_bias_into_conv_use_move_constant(self): # type: () -> None
conv = helper.make_node("Conv", ["X", "Y"], ["Z"])
constant = helper.make_node("Constant", [], ["A"],
value=helper.make_tensor(
name="bias",
data_type=TensorProto.FLOAT,
dims=(16,),
vals=np.random.randn(16).astype(np.float32).tolist()))
add = helper.make_node("Add", ["Z", "A"], ["B"])
graph = helper.make_graph(
[conv, constant, add],
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 5, 3, 3)),
helper.make_tensor_value_info("Y", TensorProto.FLOAT, (16, 5, 3, 3))],
[helper.make_tensor_value_info("B", TensorProto.FLOAT, (1, 16, 3, 3))],
value_info=[
helper.make_tensor_value_info("A", TensorProto.FLOAT, (16, 1, 1)),
]
)
optimized_model = self._optimized(graph, ["fuse_add_bias_into_conv"])
assert len(optimized_model.graph.node) == 3
assert optimized_model.graph.node[0].op_type == 'Constant'
assert optimized_model.graph.node[1].op_type == 'Squeeze'
assert optimized_model.graph.node[2].op_type == 'Conv'
assert optimized_model.graph.output[0].name == 'Z'
assert optimized_model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT
assert len(optimized_model.graph.output[0].type.tensor_type.shape.dim) == 4
def test_lift_lex_if(self): # type: () -> None
nodes = [helper.make_node("Identity", ["X"], ["Y"])]
nodes.extend(self._make_fake_if_op(
[helper.make_node("Identity", ["X"], ["_Y2"]),
helper.make_node("Identity", ["Y"], ["_Y3"])],
[helper.make_node("Identity", ["X"], ["_Y2"]),
helper.make_node("Identity", ["X"], ["_Y3"])],
[(TensorProto.FLOAT, (5,), "Y2"),
(TensorProto.FLOAT, (5,), "Y3")]))
graph = helper.make_graph(
nodes,
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,))],
[helper.make_tensor_value_info("Y", TensorProto.FLOAT, (5,)),
helper.make_tensor_value_info("Y2", TensorProto.FLOAT, (5,))])
# "If" node now diverges from ONNX schema. Disable checking.
optimized_model = self._optimized(graph, ["lift_lexical_references"])
# Identity, Constant (condition), If
assert len(optimized_model.graph.node) == 3
# else_branch, then_branch, __control_inputs
assert len(optimized_model.graph.node[2].attribute) == 3
assert optimized_model.graph.node[2].attribute[2].name == "__control_inputs"
assert optimized_model.graph.node[2].attribute[2].strings[0] == b"X"
assert optimized_model.graph.node[2].attribute[2].strings[1] == b"Y"
def test_eliminate_identity_single_use(self): # type: () -> None
nodes = [helper.make_node("Identity", ["X"], ["Y"])]
nodes.extend(self._make_fake_loop_op(
[helper.make_node("Identity", ["_Y"], ["_Y2"])],
[(TensorProto.FLOAT, (5,), "Y")],
[(TensorProto.FLOAT, (5,), "Y2")]))
graph = helper.make_graph(
nodes,
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,))],
[helper.make_tensor_value_info("Y", TensorProto.FLOAT, (5,)),
helper.make_tensor_value_info("Y2", TensorProto.FLOAT, (5,))])
optimized_model = self._optimized(graph, ["eliminate_identity"])
# All identity nodes should have been eliminated
def check_identity(node): # type: (NodeProto) -> None
assert node.op_type != "Identity"
self._visit_all_nodes_recursive(optimized_model.graph, check_identity)
# Use of the output from the Identity node in the main graph should
# have been replaced with the input to the identity node
assert len(optimized_model.graph.output) == 2
assert optimized_model.graph.output[0].name == "X"
# Use of the output from the Identity node in the loop graph should
# have been replaced with the input to that identity node
assert len(optimized_model.graph.node[2].attribute[0].g.output) == 2
assert optimized_model.graph.node[2].attribute[0].g.output[1].name == "_Y"
def _create_gemm(cls, op_def, shapes):
x, w, b = op_def.input
args = {arg.name: arg for arg in op_def.arg}
y, = op_def.output
x_shape = list(shapes[x])
nodes = []
const_tensors = []
if 'axis' in args:
axis = args['axis'].i
outer = np.prod(x_shape[:axis]).astype(int)
inner = np.prod(x_shape[axis:]).astype(int)
reshaped_x = dummy_name()
shape_tensor = cls._create_shape_tensor([outer, inner])
const_tensors.append(shape_tensor)
nodes.append(helper.make_node(
'Reshape',
inputs=[x, shape_tensor.name],
outputs=[reshaped_x],
))
x = reshaped_x
if 'axis_w' in args:
axis_w = args['axis_w'].i
w_shape = shapes[w]
outer = np.prod(w_shape[:axis_w]).astype(int).item()
inner = np.prod(w_shape[axis_w:]).astype(int).item()
reshaped_w = dummy_name()
shape_tensor = cls._create_shape_tensor([outer, inner])
const_tensors.append(shape_tensor)
nodes.append(helper.make_node(
'Reshape',
inputs=[w, shape_tensor.name],
outputs=[reshaped_w],
))
w = reshaped_w
gemm_y_output = dummy_name() if 'axis' in args else y
nodes.append(helper.make_node(
'Gemm',
inputs=[x, w, b],
outputs=[gemm_y_output],
name=op_def.name,
transB=1,
broadcast=1,
))
if 'axis' in args:
axis = args['axis'].i
shape_tensor = cls._create_shape_tensor(x_shape[:axis] + [-1])
const_tensors.append(shape_tensor)
nodes.append(helper.make_node(
'Reshape',
inputs=[gemm_y_output, shape_tensor.name],
outputs=[y],
))
return nodes, const_tensors
def test_GLU_partial(self): # type: () -> None
graph = self._make_graph(
[('x', TensorProto.FLOAT, (5, 6, 7))],
[make_node('Split', ['x'], ['y', 'z'], axis=1),
make_node('Sigmoid', ['z'], ['a'])],
[])
self._assert_inferred(graph, [make_tensor_value_info('y', TensorProto.FLOAT, (5, 3, 7)),
make_tensor_value_info('z', TensorProto.FLOAT, (5, 3, 7)),
make_tensor_value_info('a', TensorProto.FLOAT, (5, 3, 7))])
def test_fuse_transpose_default(self): # type: () -> None
trans1 = helper.make_node("Transpose", ["X"], ["Y"])
trans2 = helper.make_node("Transpose", ["Y"], ["Z"])
graph = helper.make_graph(
[trans1, trans2],
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4))],
[helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 3, 4))])
optimized_model = self._optimized(graph, ["fuse_consecutive_transposes"])
assert len(list(optimized_model.graph.node)) == 0
def test_fuse_transpose(self):
trans1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[1,0,2])
trans2 = helper.make_node("Transpose", ["Y"], ["Z"], perm=[2,0,1])
trans3 = helper.make_node("Transpose", ["Z"], ["A"], perm=[2,0,1])
graph = helper.make_graph(
[trans1, trans2, trans3],
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4))],
[helper.make_tensor_value_info("A", TensorProto.FLOAT, (4, 3, 2))])
optimized_model = self._optimized(graph)
assert len(list(optimized_model.graph.node)) == 1
def test_fuse_transpose_into_gemm(self):
trans1 = helper.make_node("Transpose", ["X"], ["A"], perm=[1,0])
trans2 = helper.make_node("Transpose", ["Y"], ["B"], perm=[1,0])
gemm = helper.make_node("Gemm", ["A", "B", "C"], ["Z"])
graph = helper.make_graph(
[trans1, trans2, gemm],
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3)),
helper.make_tensor_value_info("Y", TensorProto.FLOAT, (5, 2)),
helper.make_tensor_value_info("C", TensorProto.FLOAT, (3, 5))],
[helper.make_tensor_value_info("Z", TensorProto.FLOAT, (3, 5))])
optimized_model = self._optimized(graph)
assert len(list(optimized_model.graph.node)) == 1
def test_check_node_input_marked_optional(self): # type: () -> None
# Constant fill's input is marked optional
node = helper.make_node(
"ConstantFill", [], ["Y"], name="test")
checker.check_node(node)
# Explicitly pass the empty string as optional
node = helper.make_node(
"ConstantFill", [""], ["Y"], name="test")
# Input of RELU is not optional
node = helper.make_node(
"Relu", [""], ["Y"], name="test")
self.assertRaises(checker.ValidationError, checker.check_node, node)
def test_fuse_consecutive_squeezes_default(self): # type: () -> None
squeeze1 = helper.make_node("Squeeze", ["X"], ["Y"], axes=[0, 4, 5])
squeeze2 = helper.make_node("Squeeze", ["Y"], ["Z"], axes=[0, 3])
squeeze3 = helper.make_node("Squeeze", ["Z"], ["A"], axes=[2])
nodes = [squeeze1, squeeze2, squeeze3]
graph = helper.make_graph(
nodes,
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 1, 2, 3, 1, 1, 1, 1, 8, 9))],
[helper.make_tensor_value_info("A", TensorProto.FLOAT, (2, 3, 8, 9))])
optimized_model = self._optimized(graph, ["fuse_consecutive_squeezes"])
assert optimized_model.graph.node[0].op_type == "Squeeze"
assert list(optimized_model.graph.node[0].attribute[0].ints) == [0, 1, 4, 5, 6, 7]
assert len(list(optimized_model.graph.node)) == 1
def test_check_arguments(self):
X = np.random.randn(3, 2).astype(np.float32)
Y = np.random.randn(3, 2).astype(np.float32)
Z = np.zeros((3, 2)).astype(np.float32)
b2 = C.Caffe2Backend()
node_def = make_node("Add", inputs = ["X", "Y"], outputs = ["Z"], broadcast = 0)
output = b2.convert_node(node_def.SerializeToString(), 6)
bad_node_def = make_node("Add", inputs = ["X", "Y"], outputs = ["Z"], foo = 42, bar = 56)
with self.assertRaisesRegexp(
RuntimeError,
".*?Don't know how to map unexpected argument (foo|bar) \(from operator .*?\).*$"):
output = b2.convert_node(bad_node_def.SerializeToString(), 6)
def test_eliminate_identity_multiple_uses(self): # type: () -> None
identity = helper.make_node("Identity", ["X"], ["Y"])
add = helper.make_node("Add", ["Z", "Y"], ["A"])
mul = helper.make_node("Mul", ["A", "Y"], ["B"])
graph = helper.make_graph(
[identity, add, mul],
"test",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,)),
helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5,))],
[helper.make_tensor_value_info("B", TensorProto.FLOAT, (5,))])
optimized_model = self._optimized(graph, ["eliminate_identity"])
for node in optimized_model.graph.node:
assert node.op_type != "Identity"
assert len(optimized_model.graph.node) == 2
def test_spacetodepth():
n, c, h, w = shape = (1, 1, 4, 6)
input1 = np.random.rand(n, c, h, w).astype("float32")
blocksize = 2
inputs = [helper.make_tensor_value_info("input1", TensorProto.FLOAT, shape=shape)]
outputs = [helper.make_tensor_value_info("output", TensorProto.FLOAT, shape=(1, 4, 2, 3))]
nodes = [helper.make_node("SpaceToDepth", ["input1"], ["output"], block_size=blocksize)]
graph = helper.make_graph(nodes,
"spacetodepth_test",
inputs,
outputs)
spacetodepth_model = helper.make_model(graph)
bkd_rep = backend.prepare(spacetodepth_model)
output = bkd_rep.run([input1])
tmp = np.reshape(input1, [n, c,
h // blocksize, blocksize,
w // blocksize, blocksize])
tmp = np.transpose(tmp, [0, 3, 5, 1, 2, 4])
numpy_op = np.reshape(tmp, [n, c * (blocksize**2),
h // blocksize,
w // blocksize])
npt.assert_almost_equal(output[0], numpy_op)
def test_conv_transpose(self): # type: () -> None
graph = self._make_graph(
[('X', TensorProto.FLOAT, (25, 48, 16, 16)),
('W', TensorProto.FLOAT, (48, 32, 3, 3))],
[make_node('ConvTranspose', ['X', 'W'], 'Y', strides=[2, 2])],
[])
self._assert_inferred(graph, [make_tensor_value_info('Y', TensorProto.FLOAT, (25, 32, 33, 33))])
def test_small_model(self):
# Create one input
X = helper.make_tensor_value_info('IN', TensorProto.FLOAT, [2, 3])
# Create one output
Y = helper.make_tensor_value_info('OUT', TensorProto.FLOAT, [2, 3])
# Create a node
node_def = helper.make_node('Abs', ['IN'], ['OUT'])
# Create the model
graph_def = helper.make_graph([node_def], "test-model", [X], [Y])
onnx_model = helper.make_model(graph_def,
producer_name='onnx-example')
model = Model()
model.BuildFromOnnxModel(onnx_model)
schedule = model.OptimizeSchedule()
schedule = schedule.replace('\n', ' ')
expected_schedule = r'// Target: .+// MachineParams: .+// Delete this line if not using Generator Pipeline pipeline = get_pipeline\(\);.+Func OUT = pipeline.get_func\(1\);.+{.+}.+'
self.assertRegex(schedule, expected_schedule)
input = np.random.rand(2, 3) - 0.5
outputs = model.run([input])
self.assertEqual(1, len(outputs))
output = outputs[0]
expected = np.abs(input)
np.testing.assert_allclose(expected, output)
def test_conv_transpose_with_kernal_shape(self): # type: () -> None
graph = self._make_graph(
[('X', TensorProto.FLOAT, (25, 48, 16, 16)),
('W', TensorProto.FLOAT, (48, 32, None, None))],
[make_node('ConvTranspose', ['X', 'W'], 'Y', kernel_shape=[3, 3], strides=[2, 2], pads=[1, 1, 2, 2])],
[])
self._assert_inferred(graph, [make_tensor_value_info('Y', TensorProto.FLOAT, (25, 32, 30, 30))])
def test_scalars(self):
# Create 2 inputs
X = helper.make_tensor_value_info('A', TensorProto.INT32, [])
Y = helper.make_tensor_value_info('B', TensorProto.INT32, [])
# Create one output
Z = helper.make_tensor_value_info('C', TensorProto.INT32, [])
# Create a node
node_def = helper.make_node('Add', ['A', 'B'], ['C'])
# Create the model
graph_def = helper.make_graph([node_def], "scalar-model", [X, Y], [Z])
onnx_model = helper.make_model(graph_def,
producer_name='onnx-example')
model = Model()
model.BuildFromOnnxModel(onnx_model)
schedule = model.OptimizeSchedule()
schedule = schedule.replace('\n', ' ')
expected_schedule = r'// Target: .+// MachineParams: .+// Delete this line if not using Generator Pipeline pipeline = get_pipeline\(\);.+Func C = pipeline.get_func\(2\);.+{.+}.+'
self.assertRegex(schedule, expected_schedule)
input1 = np.random.randint(-10, 10, size=())
input2 = np.random.randint(-10, 10, size=())
outputs = model.run([input1, input2])
self.assertEqual(1, len(outputs))
output = outputs[0]
expected = input1 + input2
np.testing.assert_allclose(expected, output)
def test_split_from_GLU(self): # type: () -> None
graph = self._make_graph(
[('x', TensorProto.FLOAT, (5, 6, 7))],
[make_node('Split', ['x'], ['y', 'z'], axis=1)],
[])
self._assert_inferred(graph, [make_tensor_value_info('y', TensorProto.FLOAT, (5, 3, 7)),
make_tensor_value_info('z', TensorProto.FLOAT, (5, 3, 7))])
def test_roipool(self): # type: () -> None
graph = self._make_graph(
[("X", TensorProto.FLOAT, (5, 3, 4, 4)),
("rois", TensorProto.INT64, (2, 5))],
[make_node("MaxRoiPool", ["X", "rois"], ["Y"], pooled_shape=[2, 2])],
[])
self._assert_inferred(graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, 3, 2, 2))])
def ModelTransfer(model_path, output_path):
# 读取模型文件，该模型文件存储各层类型和名称，以及参数存储位置
with open(model_path, "r") as f:
model_define = json.load(f)
node_list = []
input_list = []
output_list = []
size = model_define["0"]["input_size"]
if len(size) > 2:
size = [size[0], size[3], size[1], size[2]]
input_list.append(
helper.make_tensor_value_info(model_define["0"]["input_name"],
TensorProto.FLOAT, size))
# 根据各层类型判断所需要生成的结点
for index in range(len(model_define)):
node = model_define[str(index)]
# 对卷积层生成一个乘法和加法
if node["type"] == "FC":
s = np.load(node["weights_path"]).astype(np.float32)
node_list.append(
helper.make_node(
"Constant",
inputs=[],
outputs=[node["weights_name"]],
value=helper.make_tensor(
name=node["weights_name"],
data_type=TensorProto.FLOAT,
dims=s.shape,
vals=s.flatten().astype(float),
),
))
s = np.load(node["bias_path"]).astype(np.float32)
node_list.append(
helper.make_node(
"Constant",
inputs=[],
outputs=[node["bias_name"]],
value=helper.make_tensor(
name=node["bias_name"],
data_type=TensorProto.FLOAT,
dims=s.shape,
vals=s.flatten().astype(float),
),
))
node_list.append(
helper.make_node("MatMul",
[node["input_name"], node["weights_name"]],
[node["output_name"] + "Temp"]))
node_list.append(
helper.make_node(
"Add", [node["output_name"] + "Temp", node["bias_name"]],
[node["output_name"]]))
# 卷积层对应代码
elif node["type"] == "Conv":
s = np.load(node["weights_path"]).astype(np.float32)
s = np.swapaxes(s, 0, 3)
s = np.swapaxes(s, 1, 2)
# s = np.swapaxes(s, 2, 3)
node_list.append(
helper.make_node(
"Constant",
inputs=[],
outputs=[node["weights_name"]],
value=helper.make_tensor(
name=node["weights_name"],
data_type=TensorProto.FLOAT,
dims=s.shape,
vals=s.flatten().astype(float),
),
))
# should use broadcast, but I didn't find how to use that attribute
s = np.load(node["bias_path"]).astype(np.float32)
s = np.tile(s, node["output_size"][1:3] + [1])
s = np.swapaxes(s, 0, 2)
s = np.swapaxes(s, 1, 2)
node_list.append(
helper.make_node(
"Constant",
inputs=[],
outputs=[node["bias_name"]],
value=helper.make_tensor(
name=node["bias_name"],
data_type=TensorProto.FLOAT,
dims=s.shape,
vals=s.flatten().astype(float),
),
))
node_list.append(
helper.make_node(node["type"],
[node["input_name"], node["weights_name"]],
[node["output_name"] + "Temp"],
kernel_shape=node["kernel_shape"],
strides=node["strides"],
pads=node["pads"]))
node_list.append(
helper.make_node(
"Add",
inputs=[node["output_name"] + "Temp", node["bias_name"]],
outputs=[node["output_name"]]))
# relu和softmax
elif node["type"] == "Relu" or node["type"] == "Softmax" or node[
"type"] == "Sigmoid" or node["type"] == "Tanh":
node_list.append(
helper.make_node(node["type"], [node["input_name"]],
[node["output_name"]]))
# max pooling
elif node["type"] == "MaxPool":
node_list.append(
helper.make_node(node["type"], [node["input_name"]],
[node["output_name"]],
kernel_shape=node["kernel_shape"],
strides=node["kernel_shape"],
pads=node["pads"]))
# reshape对应代码，添加转置，方便与numpy代码接轨
elif node["type"] == "Reshape":
shape = np.array(node["shape"], dtype=np.int64)
node_list.append(
helper.make_node('Transpose',
inputs=[node["input_name"]],
outputs=[node["input_name"] + "T"],
perm=[0, 2, 3, 1]))
node_list.append(
helper.make_node(
"Constant",
inputs=[],
outputs=[node["output_name"] + "shape"],
value=helper.make_tensor(
name=node["output_name"] + "shape",
data_type=TensorProto.INT64,
dims=shape.shape,
vals=shape.flatten(),
),
))
node_list.append(
helper.make_node(
node["type"],
[node["input_name"] + "T", node["output_name"] + "shape"],
[node["output_name"]],
))
size = model_define[str(index)]["output_size"]
if len(size) > 2:
size = [size[0], size[3], size[1], size[2]]
output_list.append(
helper.make_tensor_value_info(model_define[str(index)]["output_name"],
TensorProto.FLOAT, size))
graph_proto = helper.make_graph(
node_list,
"test",
input_list,
output_list,
)
onnx.checker.check_node(node_list[1])
onnx.checker.check_graph(graph_proto)
model_def = helper.make_model(graph_proto, producer_name="test_onnx")
onnx.checker.check_model(model_def)
onnx.save(model_def, output_path)
def fuse(self, node, input_name_to_nodes, output_name_to_node):
if self.model.match_parent_path(node, ['Add', 'Gather'],
[0, 0]) is None:
logger.debug(
"Failed to match path SkipLayerNormalization[0] <-- Add <-- Gather"
)
return
self.attention = self.model.find_first_child_by_type(
node, 'Attention', input_name_to_nodes, recursive=False)
if self.attention is None:
# In case user disables attention fusion, check whether subgraph looks like Attention.
if node.output[0] not in input_name_to_nodes:
return
children = input_name_to_nodes[node.output[0]]
children_types = sorted([child.op_type for child in children])
if children_types != [
'MatMul', 'MatMul', 'MatMul', 'SkipLayerNormalization'
]:
logger.debug(
"No Attention like subgraph in children of SkipLayerNormalization"
)
return
# Assume the order of embeddings are word_embedding + position_embedding + segment_embedding
normalize_node = node
word_embedding_path = self.model.match_parent_path(
normalize_node, ['Add', 'Gather'], [0, 0])
if word_embedding_path is None:
logger.info(
"Word embedding path is not found. Embed layer cannot be fused."
)
return
add_node, word_embedding_gather = word_embedding_path
input_ids = word_embedding_gather.input[1]
position_embedding_expand = None
position_embedding_shape = None
position_embedding_path = self.model.match_parent_path(
normalize_node, ['Reshape', 'Slice'], [1, 0])
if position_embedding_path is not None:
_, position_embedding_weight_node = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(
add_node, ['Gather', 'Expand', 'Shape'], [1, 1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand, position_embedding_shape = position_embedding_path
else:
position_embedding_path = self.model.match_parent_path(
add_node, [
'Gather', 'Expand', 'Concat', 'Unsqueeze', 'Gather',
'Shape'
], [1, 1, 1, 1, 0, 0])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand, _, _, _, position_embedding_shape = position_embedding_path
else:
# Here we will not try to get exact match. Instead, we only try identify position embedding weights.
position_embedding_path = self.model.match_parent_path(
add_node, ['Gather', 'Expand'], [1, 1])
if position_embedding_path is not None:
position_embedding_weight_node, position_embedding_expand = position_embedding_path
else:
logger.info(
"Position embedding path is not found. Embed layer cannot be fused."
)
return
if position_embedding_shape is not None and position_embedding_shape.input[
0] != input_ids:
logger.info(
"position and word embedding is expected to be applied on same input"
)
return
segment_embedding_path = self.model.match_parent_path(
normalize_node, ['Gather'], [1])
if segment_embedding_path is None:
segment_embedding_path = self.model.match_parent_path(
normalize_node, ['Add', 'Gather'], [0, 1])
if segment_embedding_path is None:
logger.info(
"Segment embedding is not found. Embed layer cannot be fused."
)
return
_, segment_embedding_gather = segment_embedding_path
else:
segment_embedding_gather = segment_embedding_path[0]
segment_ids = segment_embedding_gather.input[1]
if position_embedding_expand and position_embedding_shape:
input_parent = self.model.get_parent(position_embedding_shape, 0,
output_name_to_node)
subgraph_nodes = self.model.get_parent_subgraph_nodes(
position_embedding_expand,
[input_parent] if input_parent else [], output_name_to_node)
self.nodes_to_remove.extend(subgraph_nodes)
self.nodes_to_remove.extend(word_embedding_path)
self.nodes_to_remove.extend(position_embedding_path)
self.nodes_to_remove.extend(segment_embedding_path)
self.nodes_to_remove.extend([normalize_node])
# Cast input_ids and segment_ids to int32.
input_ids_cast_node = None
if self.model.find_graph_input(input_ids):
casted, input_ids = self.utils.cast_graph_input_to_int32(input_ids)
else:
input_ids, input_ids_cast_node = self.utils.cast_input_to_int32(
input_ids)
if self.model.find_graph_input(segment_ids):
casted, segment_ids = self.utils.cast_graph_input_to_int32(
segment_ids)
else:
segment_ids, segment_ids_cast_node = self.utils.cast_input_to_int32(
segment_ids)
# Cast might be removed by OnnxRuntime.
_, segment_id_path, _ = self.model.match_parent_paths(
segment_ids_cast_node, [([
'ConstantOfShape', 'Concat', 'Unsqueeze', 'Gather',
'Shape', 'Cast'
], [0, 0, 1, 0, 0, 0]),
([
'ConstantOfShape', 'Concat',
'Unsqueeze', 'Gather', 'Shape'
], [0, 0, 1, 0, 0])],
output_name_to_node)
if segment_id_path and input_ids_cast_node and input_ids_cast_node.input[
0] == segment_id_path[-1].input[0]:
logger.debug("Simplify semgent id path...")
self.model.add_node(
helper.make_node('Shape',
inputs=[input_ids_cast_node.input[0]],
outputs=["input_shape"]))
self.model.add_node(
helper.make_node('ConstantOfShape',
inputs=["input_shape"],
outputs=["zeros_for_input_shape"],
value=helper.make_tensor(
"value", onnx.TensorProto.INT32, [1],
[1])))
segment_ids = "zeros_for_input_shape"
node_name = self.model.create_node_name('EmbedLayerNormalization')
output_name = node_name + "_output"
embed_node = helper.make_node(
'EmbedLayerNormalization',
inputs=[
input_ids,
segment_ids,
word_embedding_gather.input[0],
position_embedding_weight_node.input[0],
segment_embedding_gather.input[0],
normalize_node.input[2],
normalize_node.input[3] # gamma and beta
],
outputs=[node_name + "_output", node_name + "_dummy_mask_index"],
name=node_name)
embed_node.domain = "com.microsoft"
# Pass attribute "epsilon" from normalize node to EmbedLayerNormalization.
for att in normalize_node.attribute:
if att.name == 'epsilon':
embed_node.attribute.extend([att])
# Set default value to 1e-12 if no attribute is found.
# OnnxRuntime 1.2.0 or older has no epsilon attribute. The optimized model can only work for 1.3.0 or later.
if len(embed_node.attribute) == 0:
embed_node.attribute.extend(
[helper.make_attribute("epsilon", 1.0E-12)])
self.model.replace_input_of_all_nodes(normalize_node.output[0],
output_name)
self.nodes_to_add.append(embed_node)
def test_leakyrelu(self):
node_def = helper.make_node("LeakyRelu", ["X"], ["Y"], alpha=2.0)
x = np.floor(self._get_rnd([100]))
output = run_node(node_def, [x])
test_output = [self._leaky_relu(a, 2.0) for a in x]
np.testing.assert_almost_equal(output["Y"], test_output)
def test_relu(self):
node_def = helper.make_node("Relu", ["X"], ["Y"])
x = self._get_rnd([1000])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.maximum(x, 0))
def test_check_graph_optional_input(self): # type: () -> None
node = helper.make_node("ConstantFill", [""], ["Y"], name="test")
graph = helper.make_graph(
[node], "test", [],
[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2])])
checker.check_graph(graph)
def fuse_attention(self):
output_name_to_node = self.output_name_to_node()
nodes_to_remove = []
attention_count = 0
start_nodes = []
skip_layer_norm_nodes = self.get_nodes_by_op_type(
"SkipLayerNormalization")
layer_norm_nodes = self.get_nodes_by_op_type("LayerNormalization")
# Sometimes we can not fuse skiplayernormalization since the add before layernorm has an output that used by nodes outside skiplayernorm
# Conceptually we treat add before layernorm as skiplayernorm node since they share the same pattern
start_nodes.extend(skip_layer_norm_nodes)
start_nodes.extend(layer_norm_nodes)
graph_name = self.get_graph_by_node(start_nodes[0]).name
for normalize_node in start_nodes:
# SkipLayerNormalization has two inputs, and one of them is the root input for attention.
if normalize_node.op_type == 'LayerNormalization':
add_before_layernorm = self.match_parent(
normalize_node, 'Add', 0)
if add_before_layernorm is not None:
normalize_node = add_before_layernorm
else:
continue
parent = self.get_parent(normalize_node, 1)
if parent is None or parent.op_type not in [
"SkipLayerNormalization", "LayerNormalization", "Reshape"
]:
parent = self.get_parent(normalize_node, 0)
if parent is None or parent.op_type not in [
"SkipLayerNormalization", "LayerNormalization",
"Reshape"
]:
logger.debug("Failed to match parent of normalize_node")
continue
qkv_nodes = self.match_parent_path(
normalize_node,
['Add', 'MatMul', 'Reshape', 'Transpose', 'MatMul'],
[0, 0, 0, 0, 0])
if qkv_nodes is None:
qkv_nodes = self.match_parent_path(
normalize_node,
['MatMul', 'Reshape', 'Transpose', 'MatMul'], [1, 0, 0, 0])
if qkv_nodes is None:
qkv_nodes = self.match_parent_path(
normalize_node, ['Add', 'Einsum', 'Einsum'], [0, 0, 0])
if qkv_nodes is None:
logger.debug("Failed to match qkv nodes")
continue
matmul_qkv = qkv_nodes[-1]
v_nodes = self.match_parent_path(
matmul_qkv, ['Transpose', 'Reshape', 'Add', 'MatMul'],
[1, 0, 0, 0])
if v_nodes is None:
v_nodes = self.match_parent_path(matmul_qkv, ['Add', 'Einsum'],
[1, 0])
if v_nodes is None:
logger.debug("Failed to match v path")
continue
add_v = v_nodes[-2]
matmul_v = v_nodes[-1]
qk_nodes = self.match_parent_path(
matmul_qkv, ['Softmax', 'Add', "Mul", 'MatMul'], [0, 0, 0, 0])
if qk_nodes is None:
qk_nodes = self.match_parent_path(matmul_qkv,
['Softmax', 'Add', 'Einsum'],
[0, 0, 0])
if qk_nodes is None:
logger.debug("Failed to match qk_paths")
continue
matmul_qk = qk_nodes[-1]
q_nodes = self.match_parent_path(
matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'],
[0, 0, 0, 0])
if q_nodes is None:
q_nodes = self.match_parent_path(matmul_qk, ['Add', 'Einsum'],
[0, 0])
if q_nodes is None:
logger.debug("Failed to match q path")
continue
add_q = q_nodes[-2]
matmul_q = q_nodes[-1]
k_nodes = self.match_parent_path(
matmul_qk, ['Transpose', 'Reshape', 'Add', 'MatMul'],
[1, 0, 0, 0])
if k_nodes is None:
k_nodes = self.match_parent_path(matmul_qk,
['Mul', 'Add', 'Einsum'],
[1, 0, 0])
if k_nodes is None:
logger.debug("Failed to match k path")
continue
add_k = k_nodes[-2]
matmul_k = k_nodes[-1]
mask_nodes = self.match_mask_path(qk_nodes[1])
if mask_nodes is None:
logger.debug("Cannot find mask_nodes.")
continue
if not self.has_constant_input(mask_nodes[1], 1):
logger.debug(
"Sub node expected to have an input with constant value 1.0."
)
continue
# add a squeeze node to convert a 3-d mask to 2-d
squeeze_node = self.match_parent_path(
mask_nodes[-1], ['Squeeze'], [0]) or self.match_parent_path(
mask_nodes[-1], ['Expand'], [0])
squeeze_node_name = "Squeeze_3d_to_2d_mask"
squeeze_output_name = squeeze_node_name + "_output"
if squeeze_node is None and len(
mask_nodes) == 5 and self.find_graph_input(
mask_nodes[-1].input[0]) is None:
mask_input = mask_nodes[-1].input[1]
self.add_node(
helper.make_node("Squeeze", [mask_input],
[squeeze_output_name],
squeeze_node_name,
axes=[1]), graph_name)
mask_nodes[-1].input[0] = squeeze_output_name
is_same_root = self.check_attention_input(matmul_q, matmul_k,
matmul_v, parent,
output_name_to_node)
if is_same_root:
mask_index = self.attention_mask.process_mask(
mask_nodes[-1].input[0])
logger.debug("Create an Attention node.")
# For tf models, q and v are flipped.
attention_node = self.attention_fusion.create_attention_node(
mask_index, matmul_k, matmul_q, matmul_v, add_k, add_q,
add_v, self.num_heads, self.hidden_size, parent.output[0],
qkv_nodes[2].output[0])
if attention_node is None:
continue
if qkv_nodes[1].op_type == 'Einsum':
# add reshape before einsum
tensor = helper.make_tensor(
name=qkv_nodes[1].name + "_newshape",
data_type=TensorProto.INT64,
dims=[4],
vals=np.int64([[
0, 0, self.num_heads,
int(self.hidden_size / self.num_heads)
]]).tobytes(),
raw=True)
self.add_initializer(tensor, graph_name)
reshape_ = helper.make_node(
"Reshape",
inputs=[
attention_node.output[0],
qkv_nodes[1].name + "_newshape"
],
outputs=[qkv_nodes[1].name + "_reshape_output"],
name=qkv_nodes[1].name + "_reshape")
qkv_nodes[1].input[
0] = qkv_nodes[1].name + "_reshape_output"
self.add_node(reshape_, graph_name)
if parent.op_type == 'Reshape':
# Temporary work around: we require the skiplayernorm and attention op be fed with 3-d input
hidden_size = numpy_helper.to_array(
self.get_initializer(parent.input[1]))[1]
tensor = helper.make_tensor(
name=parent.name + "_modified",
data_type=TensorProto.INT64,
dims=[3],
vals=np.int64([[1, -1, hidden_size]]).tobytes(),
raw=True)
self.add_initializer(tensor, graph_name)
parent.input[1] = parent.name + "_modified"
self.add_node(attention_node, graph_name)
attention_count += 1
nodes_to_remove.extend(qkv_nodes[2:])
nodes_to_remove.extend(qk_nodes)
nodes_to_remove.extend(q_nodes)
nodes_to_remove.extend(k_nodes)
nodes_to_remove.extend(v_nodes)
nodes_to_remove.extend(mask_nodes)
else:
logger.debug("Root node not matched.")
continue
self.remove_nodes(nodes_to_remove)
self.update_graph()
logger.info(f"Fused Attention count:{attention_count}")
def test_transpose(self):
node_def = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 1])
x = self._get_rnd([1000]).reshape([10, 10, 10])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.transpose(x, (0, 2, 1)))
def test_exp(self):
node_def = helper.make_node("Exp", ["X"], ["Y"])
x = self._get_rnd([100])
x = x - 3.6
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.exp(x))
def test_reduce_l1(self):
node_def = helper.make_node("ReduceL1", ["X"], ["Y"], axes=[1, 2])
x = self._get_rnd([5, 10, 10, 3])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"],
np.linalg.norm(x, 1, (1, 2), True))
def test_sqrt(self):
node_def = helper.make_node("Sqrt", ["X"], ["Y"])
x = self._get_rnd([1000]) + 1.0
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.sqrt(x), decimal=5)
def test_sigmoid(self):
node_def = helper.make_node("Sigmoid", ["X"], ["Y"])
x = self._get_rnd([1000])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], 1/(1 + np.exp(-x)))
def test_reshape(self):
node_def = helper.make_node("Reshape", ["X"], ["Y"], shape=[10, 10])
x = self._get_rnd(100)
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], x.reshape([10, 10]))
def test_elu(self):
node_def = helper.make_node("Elu", ["X"], ["Y"])
x = self._get_rnd([100])
output = run_node(node_def, [x])
test_output = [self._elu(a) for a in x];
np.testing.assert_almost_equal(output["Y"], test_output)
def handle_reshape(cls, node, **kwargs):
return helper.make_node("Reshape", [node.inputs[0], node.inputs[1]],
[node.name])
def create_loop_body_graph(gather_input_ids, output_data_type, output_shape,
trip_count_input_ids, rank, loop_name):
nodes = []
iter_name = utils.make_name("i")
cond_name = utils.make_name("cond")
fake_var_name = utils.make_name("fake_var")
graph_inputs = [
helper.make_tensor_value_info(iter_name, TensorProto.INT64,
(1, )), # iteration_num
helper.make_tensor_value_info(cond_name, TensorProto.BOOL,
()), # condition
helper.make_tensor_value_info(fake_var_name, TensorProto.FLOAT,
()) # loop-carried dependency
]
# get the i'th value of condition
cond_input_id = gather_input_ids[0]
new_nodes, cond_input_id_for_current_iter = get_inputs_for_current_iteration(
cond_input_id, iter_name)
nodes.extend(new_nodes)
# get the i'th value of true values
true_input_id = gather_input_ids[1]
new_nodes, true_input_id_for_current_iter = get_inputs_for_current_iteration(
true_input_id, iter_name)
nodes.extend(new_nodes)
# get the i'th value of false values
false_input_id = gather_input_ids[2]
new_nodes, false_input_id_for_current_iter = get_inputs_for_current_iteration(
false_input_id, iter_name)
nodes.extend(new_nodes)
input_ids_for_current_iter = [
cond_input_id_for_current_iter, true_input_id_for_current_iter,
false_input_id_for_current_iter
]
output_id = None
rank = rank - 1
if rank >= 1:
nodes_1 = create_loop_op(input_ids_for_current_iter, output_data_type,
output_shape[1:], trip_count_input_ids, rank)
loop_1 = nodes_1[-1]
output_id = loop_1.output[1]
nodes.extend(nodes_1)
elif rank == 0:
if_node, if_node_output_id = create_if_op(input_ids_for_current_iter,
output_data_type,
output_shape[1:])
output_id = if_node_output_id
nodes.append(if_node)
output_identity_name = utils.make_name("loop_output")
loop_output_id = utils.port_name(output_identity_name)
loop_output_node = helper.make_node('Identity', [output_id],
[loop_output_id],
name=output_identity_name)
nodes.append(loop_output_node)
cond_identity_name = utils.make_name("cond_output")
cond_output_id = utils.port_name(cond_identity_name)
identity_node = helper.make_node('Identity', [cond_name], [cond_output_id],
name=cond_identity_name)
nodes.append(identity_node)
fake_var_identity_name = utils.make_name("fake_var_output")
fake_var_output_id = utils.port_name(fake_var_identity_name)
identity_node = helper.make_node('Identity', [fake_var_name],
[fake_var_output_id],
name=fake_var_identity_name)
nodes.append(identity_node)
graph_outputs = [
helper.make_tensor_value_info(cond_output_id, TensorProto.BOOL, ()),
helper.make_tensor_value_info(fake_var_output_id, TensorProto.FLOAT,
()),
helper.make_tensor_value_info(loop_output_id, output_data_type,
output_shape[1:])
]
body_graph = helper.make_graph(nodes,
utils.make_name(loop_name + "-body-graph"),
graph_inputs, graph_outputs)
return body_graph
def test_sub(self):
node_def = helper.make_node("Sub", ["X", "Y"], ["Z"])
x = self._get_rnd([10, 10])
y = self._get_rnd([10, 10])
output = run_node(node_def, [x, y])
np.testing.assert_almost_equal(output["Z"], np.subtract(x, y))
def set_up_reference_model(act, idt, wdt, k, ifm_dim, ifm_ch, stride, padding):
# set up reference model consisting of Im2Col + MatMul (+ MultiThreshold)
ofm_ch = ifm_ch
total_pad = 2 * padding
ofm_dim = compute_conv_output_dim(ifm_dim, k, stride, total_pad=total_pad)
if act is None:
odt = DataType.INT32
else:
odt = act
out_act = oh.make_tensor_value_info("out_act", TensorProto.FLOAT,
[1, ofm_dim, ofm_dim, ofm_ch])
T = oh.make_tensor_value_info("T", TensorProto.FLOAT, [ofm_ch, 15])
tdt = DataType.INT32
thresh_node = oh.make_node(
"MultiThreshold",
domain="finn.custom_op.general",
inputs=["outp", "T"],
outputs=["out_act"],
data_layout="NHWC",
out_dtype=odt.name,
out_scale=1.0,
out_bias=0.0,
)
# set up onnx model
inp = oh.make_tensor_value_info("inp", TensorProto.FLOAT,
[1, ifm_dim, ifm_dim, ifm_ch])
outp = oh.make_tensor_value_info("outp", TensorProto.FLOAT,
[1, ofm_dim, ofm_dim, ofm_ch])
W_sparse = oh.make_tensor_value_info("W_sparse", TensorProto.FLOAT,
[ifm_ch * k * k, ofm_ch])
im2col_node = oh.make_node(
"Im2Col",
domain="finn.custom_op.general",
inputs=["inp"],
outputs=["im2col_out"],
kernel_size=[k, k],
stride=stride,
pad_amount=[padding, padding, padding, padding],
input_shape="(1, {}, {}, {})".format(ifm_dim, ifm_dim, ifm_ch),
depthwise=1,
)
matmul_node = oh.make_node("MatMul",
inputs=["im2col_out", "W_sparse"],
outputs=["outp"])
if act is None:
node_list = [im2col_node, matmul_node]
global_out = outp
value_info = [W_sparse]
else:
node_list = [im2col_node, matmul_node, thresh_node]
global_out = out_act
value_info = [W_sparse, T]
graph = oh.make_graph(
nodes=node_list,
name="lowered_dw_cnv_graph",
inputs=[inp],
outputs=[global_out],
value_info=value_info,
)
model = oh.make_model(graph, producer_name="lowered_dw_cnv-model")
model = ModelWrapper(model)
# initialize model
model.set_tensor_datatype("inp", idt)
model.set_tensor_datatype(model.graph.output[0].name, odt)
model.set_tensor_datatype("W_sparse", wdt)
w_tensor = gen_finn_dt_tensor(wdt, [ofm_ch, 1, k, k])
# create sparse matrix
W_matrix = np.zeros((ofm_ch, ifm_ch, k, k))
for ch in range(ifm_ch):
W_matrix[ch][ch] = w_tensor[ch][0]
W_matrix = W_matrix.astype(np.float32)
W_matrix = W_matrix.transpose(0, 2, 3, 1)
W_matrix = W_matrix.reshape(ofm_ch, ifm_ch * k * k)
model.set_initializer("W_sparse", W_matrix.T)
sparsity = {"dw": {"kernel_shape": k}}
model.set_tensor_sparsity("W_sparse", sparsity)
if act is not None:
(min, max) = calculate_signed_dot_prod_range(idt, wdt, ifm_ch * k * k)
n_steps = odt.get_num_possible_values() - 1
T_values = np.random.randint(min, max - 1,
(ofm_ch, n_steps)).astype(np.float32)
# provide non-decreasing thresholds
T_values = np.sort(T_values, axis=1)
model.set_initializer("T", T_values)
model.set_tensor_datatype("T", tdt)
model = model.transform(InferShapes())
return model
def test_div(self):
node_def = helper.make_node("Div", ["X", "Y"], ["Z"], broadcast=1)
x = self._get_rnd([10, 10])
y = self._get_rnd([10, 10])
output = run_node(node_def, [x, y])
np.testing.assert_almost_equal(output["Z"], np.divide(x, y))
def test_check_node(self): # type: () -> None
node = helper.make_node("Relu", ["X"], ["Y"], name="test")
checker.check_node(node)
def test_mul(self):
node_def = helper.make_node("Mul", ["X", "Y"], ["Z"], broadcast=1, axis=1)
x = self._get_rnd([5, 10, 5, 5])
y = self._get_rnd([10])
output = run_node(node_def, [x, y])
np.testing.assert_almost_equal(output["Z"], np.multiply(x, y.reshape([1, 10, 1, 1])))
def test_size(self):
node_def = helper.make_node("Size", ["X"], ["Y"])
x = self._get_rnd([5, 10, 10, 3])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.size(x))
def test_neg(self):
node_def = helper.make_node("Neg", ["X"], ["Y"])
x = self._get_rnd([1000])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.negative(x))
def test_shape(self):
node_def = helper.make_node("Shape", ["X"], ["Y"])
x = self._get_rnd([5, 10, 10, 3])
output = run_node(node_def, [x])
np.testing.assert_allclose(output["Y"], np.shape(x))
def test_log(self):
node_def = helper.make_node("Log", ["X"], ["Y"])
x = self._get_rnd([100])
x = x + 3.6;
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.log(x))
def test_softsign(self):
node_def = helper.make_node("Softsign", ["X"], ["Y"])
x = self._get_rnd([3, 4, 5])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], x / (1 + np.abs(x)))
def test_softplus(self):
node_def = helper.make_node("Softplus", ["X"], ["Y"])
x = self._get_rnd([3, 4, 5])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.log(np.exp(x) + 1))
def test_squeeze(self):
node_def = helper.make_node("Squeeze", ["X"], ["Y"], axes=[2])
x = np.array([[[0], [1], [2]]])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.squeeze(x, axis=2))
def test_floor(self):
node_def = helper.make_node("Floor", ["X"], ["Y"])
x = self._get_rnd([100])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], np.floor(x))
def test_reciprocal(self):
node_def = helper.make_node("Reciprocal", ["X"], ["Y"])
x = self._get_rnd([1000])
output = run_node(node_def, [x])
np.testing.assert_almost_equal(output["Y"], 1.0/x)
def test_pow(self):
node_def = helper.make_node("Pow", ["X", "Y"], ["Z"])
x = self._get_rnd(1000) / 2.0 + 0.5
y = self._get_rnd(1000) / 2.0 + 0.5
output = run_node(node_def, [x, y])
np.testing.assert_almost_equal(output["Z"], np.power(x, y))
