def test_single_input_to_single_operation(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
x = mb.square(x=x)
return x
self.assertEqual(get_op_types_in_program(prog), ['square'])
apply_pass_and_basic_check(
prog, transform.FP16ComputePrecision(op_selector=lambda op: True))
_, _, block = apply_pass_and_basic_check(
prog, "common::dead_code_elimination")
self.assertEqual(get_op_types_in_program(prog),
["cast", "square", "cast"])
# Asserting first cast configuration
cast_1 = block.find_ops(op_type="cast")[0]
self.assertEqual(cast_1.dtype.val, "fp16")
self.assertEqual(len(cast_1.outputs), 1)
self.assertEqual(len(cast_1.outputs[0].child_ops), 1)
self.assertEqual(cast_1.outputs[0].child_ops[0].op_type, "square")
# Asserting second cast configuration
cast_2 = block.find_ops(op_type="cast")[1]
self.assertEqual(cast_2.dtype.val, "fp32")
self.assertEqual(len(cast_2.outputs), 1)
self.assertEqual(len(cast_2.outputs[0].child_ops), 0)
assert_model_is_valid(
prog,
{"x": (10, 20)},
expected_output_shapes={block.outputs[0].name: (10, 20)},
)
def test_move_split_to_first_use(self):
'''
Input graph:
x (input) ---> split ---> square ---> add (output)
| | |
| | --------------------|
|
| -----------> square --------------> relu (output)
'''
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
s1, s2 = mb.split(x=x, num_splits=2, axis=0)
x2 = mb.square(x=x)
x3 = mb.relu(x=x2)
s1_1 = mb.square(x=s1)
s3 = mb.add(x=s1_1, y=s2)
return x3, s3
assert get_op_types_in_program(prog) == ['split', 'square', 'relu', 'square', 'add']
block = prog.functions["main"]
# Reorder `split` op to test op with multiple output case
from .topological_reorder import move_operations_to_the_end_block
move_operations_to_the_end_block(block, ['split'])
assert get_op_types_in_program(prog) == ['square', 'relu', 'split', 'square', 'add']
assert_model_is_valid(
prog,
{"x": (10, 20)},
expected_output_shapes={
block.outputs[0].name: (10, 20),
block.outputs[1].name: (5, 20),
},
)
def test_0(self, reverse_order):
x_shape = tuple(np.random.randint(low=1, high=4, size=5))
@mb.program(input_specs=[mb.TensorSpec(shape=x_shape)])
def program(x):
sigmoid_x = mb.sigmoid(x=x)
if not reverse_order:
x = mb.mul(x=x, y=sigmoid_x)
else:
x = mb.mul(x=sigmoid_x, y=x)
return x
prev_prog, prev_block, block = apply_pass_and_basic_check(
program, "mil_backend::fuse_activation_silu"
)
assert get_op_types_in_program(prev_prog) == ["sigmoid", "mul"]
assert get_op_types_in_program(program) == ["silu"]
assert_model_is_valid(
program=program,
inputs={"x": x_shape},
backend=("mlprogram", "fp32"),
expected_output_shapes={block.outputs[0].name: tuple(x_shape)},
)
def test_invalid_leaky_relu_pattern2(self):
"""
Invalid because input to the "maximum" op is not same as the input of the "mul" op
Input graph:
const (val = 0.3)
|
input ----> mul ---------------> maximum -----------> output
|
const
Output graph: same as input graph
"""
@mb.program(input_specs=[mb.TensorSpec(shape=(3, 5, 6))])
def prog(x):
x1 = mb.mul(x=x, y=0.3)
x1 = mb.maximum(x=x1, y=0.4)
return x1
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::fuse_leaky_relu"
)
assert get_op_types_in_program(prev_prog) == ["mul", "maximum"]
assert get_op_types_in_program(prog) == ["mul", "maximum"]
def test_move_multiple_uses_overlapping(self):
'''
Input graph:
x (input) ---> cast ---> cast (output)
|
|-------> transpose ---> transpose (output)
'''
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
x1 = mb.cast(x=x, dtype="fp16")
x2 = mb.cast(x=x1, dtype="fp32")
x3 = mb.transpose(x=x1, perm=[1, 0])
x4 = mb.transpose(x=x3, perm=[1, 0])
return x2, x4
assert get_op_types_in_program(prog) == ['cast', 'cast', 'transpose', 'transpose']
apply_pass_and_basic_check(prog, "common::topological_reorder")
_, _, block = apply_pass_and_basic_check(prog, "common::dead_code_elimination")
assert get_op_types_in_program(prog) == ['cast', 'transpose', 'transpose', 'cast']
assert_model_is_valid(
prog,
{"x": (10, 20)},
expected_output_shapes={
block.outputs[0].name: (10, 20),
block.outputs[1].name: (10, 20)
},
)
def test_invalid_leaky_relu_pattern1(self):
"""
Invalid because alpha value greater than 1
Input graph:
const (val = 1.3)
|
input ----> mul ---------------> maximum -----------> output
| |
|----------------------------------
Output graph: same as input graph
"""
@mb.program(input_specs=[mb.TensorSpec(shape=(3, 5, 6))])
def prog(x):
x1 = mb.mul(x=x, y=1.3)
x1 = mb.maximum(x=x1, y=x)
return x1
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::fuse_leaky_relu"
)
assert get_op_types_in_program(prev_prog) == ["mul", "maximum"]
assert get_op_types_in_program(prog) == ["mul", "maximum"]
def test(self, reverse_order, elem_op):
x_shape = [1,]
@mb.program(input_specs=[mb.TensorSpec(shape=x_shape)])
def program(x):
x = mb.slice_by_index(x=x, begin=[0], end=[1], squeeze_mask=[True])
func = getattr(mb, elem_op)
if reverse_order:
x = func(x=2.0, y=x)
else:
x = func(x=x, y=2.0)
expand = mb.expand_dims(x=x, axes=[0])
other_1 = mb.add(x=x, y=[1, 2, 3])
other_2 = mb.sub(x=x, y=[1, 2, 3])
return expand, other_1, other_2
prev_prog, prev_block, block = apply_pass_and_basic_check(
program, "common::rank0_expand_dims_swap"
)
assert get_op_types_in_program(prev_prog) == ["slice_by_index", elem_op, "expand_dims", "add", "sub"]
assert get_op_types_in_program(program) == ["slice_by_index", "expand_dims", "expand_dims", elem_op, "squeeze", "add", "sub"]
assert_model_is_valid(
program=program,
inputs={"x": x_shape},
expected_output_shapes={
block.outputs[0].name: tuple(x_shape),
block.outputs[1].name: (3,),
block.outputs[2].name: (3,),
},
)
def test_linear_bias_fusion(self, rank, op_type, is_first_input, broadcast,
backend):
"""
Input graph:
Const
|
V
input -----> linear -----> add/sub ---> out
Output graph:
input -----> linear ----> out
"""
input_shape = [1, 2, 3]
input_shape = input_shape[-rank:]
input_shape = tuple(input_shape)
@mb.program(input_specs=[mb.TensorSpec(shape=input_shape)])
def prog(x):
linear_weight = np.reshape(np.arange(6), (2, 3)).astype(np.float32)
linear_bias = np.array([1., 2.])
bias = np.array([3., 4.])
if broadcast:
if rank >= 2:
bias = np.reshape(bias, (1, 2))
x = mb.linear(
x=x,
weight=linear_weight,
bias=linear_bias,
)
func = mb.add if op_type == "add" else mb.sub
if is_first_input:
kwargs = {
"x": x,
"y": bias,
}
else:
kwargs = {
"x": bias,
"y": x,
}
x = func(**kwargs)
return x
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::fuse_linear_bias")
assert get_op_types_in_program(prev_prog) == ["linear", op_type]
assert get_op_types_in_program(prog) == ["linear"]
# validate graph pass
output_shape = [1, 2, 2]
output_shape = tuple(output_shape[-rank:])
assert_model_is_valid(
prog,
{"x": input_shape},
expected_output_shapes={block.outputs[0].name: output_shape},
backend=backend,
)
def test_negative_3(self):
"""
Input graph:
input1(1, 5, 3, 4) -----> stack(axis=1) -----> reshape(shape=(-1, 2, 5, 4, 3)) ---> out(1, 5, 6, 4)
^
|
input2(1, 5, 3, 4) ----------
Output graph:
Unchanged -- this graph is not equivalent to a concat.
"""
@mb.program(input_specs=[mb.TensorSpec(shape=(1, 5, 3, 4)), mb.TensorSpec(shape=(1, 5, 3, 4))])
def prog(x1, x2):
a = mb.stack(values=[x1, x2], axis=1)
a = mb.reshape(x=a, shape=[-1, 2, 5, 4, 3])
return a
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::replace_stack_reshape"
)
self.assertEqual(
get_op_types_in_program(prev_prog), ["stack", "reshape"]
)
self.assertEqual(get_op_types_in_program(prog), ["stack", "reshape"])
def test_linear_multiple_consecutive_cast_ops(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
x = mb.cast(x=x, dtype="fp16")
x = mb.cast(x=x, dtype="fp16")
x = mb.cast(x=x, dtype="int32")
x = mb.cast(x=x, dtype="int64")
x = mb.cast(x=x, dtype="fp32")
x = mb.cast(x=x, dtype="fp16")
x = mb.square(x=x)
return x
self.assertEqual(
get_op_types_in_program(prog),
['cast', 'cast', 'cast', 'cast', 'cast', 'cast', 'square'])
apply_pass_and_basic_check(prog, "common::cast_optimization")
_, _, block = apply_pass_and_basic_check(
prog, "common::dead_code_elimination")
self.assertEqual(get_op_types_in_program(prog), ["cast", "square"])
self.assertEqual(block.find_ops(op_type="cast")[0].dtype.val, "fp16")
assert_model_is_valid(
prog,
{"x": (10, 20)},
expected_output_shapes={block.outputs[0].name: (10, 20)},
)
def test_success_3_layers(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 6, 8))])
def prog(x1):
pad1 = mb.pad(x=x1,
pad=[0, 0, 1, 1],
mode='constant',
constant_val=3.0)
pad2 = mb.pad(x=pad1,
pad=[1, 1, 0, 0],
mode='constant',
constant_val=3.0)
pad3 = mb.pad(x=pad2,
pad=[1, 1, 0, 0],
mode='constant',
constant_val=3.0)
return pad3
prev_prog, _, block = apply_pass_and_basic_check(
prog, "common::merge_consecutive_paddings")
assert get_op_types_in_program(prev_prog) == ["pad", "pad", "pad"]
assert get_op_types_in_program(prog) == ["pad"]
pad_ops = [op for op in prog["main"].operations if op.op_type == "pad"]
assert pad_ops[0].inputs["constant_val"].val == 3.0
inputs = {"x1": (1, 2, 6, 8)}
assert_model_is_valid(
prog,
inputs,
expected_output_shapes={block.outputs[0].name: (1, 2, 10, 10)},
)
def test_input_name_shadow(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
# op name "x" results in output var name "x", which shadows prog
# input var name "x"
x = mb.transpose(x=x, perm=[1, 0], name="x")
x = mb.relu(x=x, name="relu")
return x
prev_prog, _, block = apply_pass_and_basic_check(
prog, "common::dedup_op_and_var_names")
self.assertEqual(get_op_types_in_program(prev_prog),
['transpose', 'relu'])
self.assertEqual(get_op_names_in_program(prev_prog), ['x', 'relu'])
self.assertEqual(get_op_types_in_program(prog), ['transpose', 'relu'])
self.assertEqual(get_op_names_in_program(prog), ['x', 'relu'])
op = prog['main'].find_ops(op_type='transpose')[0]
self.assertEqual("x_1", op.outputs[0].name)
assert_model_is_valid(
prog,
{"x": (10, 20)},
expected_output_shapes={block.outputs[0].name: (20, 10)},
)
def test_mul_add_fusion_to_batchnorm(self, flip_mul_input_order,
flip_add_input_order,
rank_3_const_input):
C = 3
gamma = np.random.rand(1, C, 1, 1)
beta = np.random.rand(1, C, 1, 1)
if rank_3_const_input:
gamma = np.squeeze(gamma, axis=0)
beta = np.squeeze(beta, axis=0)
@mb.program(input_specs=[mb.TensorSpec(shape=(1, 10, 10, C))])
def prog(x):
x = mb.transpose(x=x, perm=[0, 3, 1, 2])
if flip_mul_input_order:
x = mb.mul(x=gamma, y=x)
else:
x = mb.mul(x=x, y=gamma)
if flip_add_input_order:
x = mb.add(x=beta, y=x)
else:
x = mb.add(x=x, y=beta)
return x
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::fuse_elementwise_to_batchnorm")
if get_op_types_in_program(prev_prog) != ["transpose", "mul", "add"]:
raise AssertionError
if get_op_types_in_program(prog) != ["transpose", "batch_norm"]:
raise AssertionError
assert_model_is_valid(
prog,
{"x": (1, 10, 10, C)},
expected_output_shapes={block.outputs[0].name: (1, C, 10, 10)},
)
def test_consecutive_fusable_casts_on_all_branches(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
x = mb.cast(x=x, dtype="int32")
x1 = mb.cast(x=x, dtype="fp16")
x2 = mb.cast(x=x, dtype="fp16")
x3 = mb.cast(x=x, dtype="fp16")
x4 = mb.square(x=x1)
x5 = mb.relu(x=x2)
x6 = mb.log(x=x3)
return x4, x5, x6
self.assertEqual(
get_op_types_in_program(prog),
['cast', 'cast', 'cast', 'cast', 'square', 'relu', 'log'])
apply_pass_and_basic_check(prog, "common::cast_optimization")
_, _, block = apply_pass_and_basic_check(
prog, "common::dead_code_elimination")
self.assertEqual(get_op_types_in_program(prog),
["cast", "square", "relu", "log"])
self.assertEqual(block.find_ops(op_type="cast")[0].dtype.val, "fp16")
assert_model_is_valid(
prog,
{"x": (10, 20)},
expected_output_shapes={
block.outputs[0].name: (10, 20),
block.outputs[1].name: (10, 20),
block.outputs[2].name: (10, 20),
},
)
def test_success_h_axis(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 6, 8))])
def prog(x1):
left = mb.slice_by_index(x=x1,
begin=[0, 0, 1, 0],
end=[0, 0, 2, 0],
end_mask=[True, True, False, True])
right = mb.slice_by_index(x=x1,
begin=[0, 0, -2, 0],
end=[0, 0, -1, 0],
end_mask=[True, True, False, True])
x = mb.concat(values=[left, x1, right], axis=2)
return x
prev_prog, _, block = apply_pass_and_basic_check(
prog, "common::use_reflection_padding")
assert get_op_types_in_program(prev_prog) == [
"slice_by_index", "slice_by_index", "concat"
]
assert get_op_types_in_program(prog) == ["pad"]
inputs = {"x1": (1, 2, 6, 8)}
assert_model_is_valid(
prog,
inputs,
expected_output_shapes={block.outputs[0].name: (1, 2, 8, 8)},
)
def test_concat_interleave_fusion_pass():
"""
Given:
%3 = concat(%1.a, %1.b, axis=-3, interleave=False) #shape = (B, n*C, H, W)
%4 = reshape(%3) #shape = (B, n, C, H, W)
%5 = transpose(%4, perm=[0, 2, 1, 3, 4]) # shape = (B, C, n, H, W)
%6 = reshape(%5) # shape = (B, C*n, H, W)
Result:
%6 = concat(%1.a, %1.b, axis=-3, interleave=True)
"""
B, C, H, W = 1, 10, 20, 20
@mb.program(input_specs=[mb.TensorSpec(shape=(B,C,H,W)), mb.TensorSpec(shape=(B,C,H,W))])
def prog(x, y):
z = mb.concat(values=[x,y], axis=1)
z = mb.reshape(x=z, shape=(B, 2, C, H, W))
z = mb.transpose(x=z, perm=[0, 2, 1, 3, 4])
z = mb.reshape(x=z, shape=(B, -1, H, W))
return z
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::detect_concat_interleave"
)
assert get_op_types_in_program(prev_prog) == ["concat", "reshape", "transpose", "reshape"]
assert get_op_types_in_program(prog) == ["concat"]
concat_op = prog.find_ops(op_type="concat", exactly_one=True)[0]
assert concat_op.interleave.val
assert_model_is_valid(
prog,
{"x": (B, C, H, W), "y": (B, C, H, W)},
expected_output_shapes={block.outputs[0].name: (B, 2*C, H, W)},
)
def test_op_name_duplicated_many(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
def prog(x):
x = mb.cast(x=x, dtype="fp16", name="castop")
x = mb.cast(x=x, dtype="fp16", name="castop")
x = mb.cast(x=x, dtype="int32", name="castop_2")
x = mb.cast(x=x, dtype="int64", name="castop")
x = mb.cast(x=x, dtype="fp32", name="castop_2")
x = mb.square(x=x, name="square")
return x
prev_prog, _, block = apply_pass_and_basic_check(
prog, "common::dedup_op_and_var_names")
self.assertEqual(get_op_types_in_program(prev_prog),
['cast', 'cast', 'cast', 'cast', 'cast', 'square'])
self.assertEqual(
get_op_names_in_program(prev_prog),
['castop', 'castop', 'castop_2', 'castop', 'castop_2', 'square'])
self.assertEqual(get_op_types_in_program(prog),
['cast', 'cast', 'cast', 'cast', 'cast', 'square'])
self.assertEqual(get_op_names_in_program(prog), [
'castop', 'castop_1', 'castop_2', 'castop_3', 'castop_2_1',
'square'
])
assert_model_is_valid(
prog,
{"x": (10, 20)},
expected_output_shapes={block.outputs[0].name: (10, 20)},
)
def test_failure_different_constants(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 6, 8))])
def prog(x1):
pad1 = mb.pad(x=x1,
pad=[0, 0, 1, 1],
mode='constant',
constant_val=1.0)
pad2 = mb.pad(x=pad1,
pad=[1, 1, 0, 0],
mode='constant',
constant_val=2.0)
return pad2
prev_prog, _, block = apply_pass_and_basic_check(
prog, "common::merge_consecutive_paddings")
assert get_op_types_in_program(prev_prog) == ["pad", "pad"]
assert get_op_types_in_program(prog) == ["pad", "pad"]
inputs = {"x1": (1, 2, 6, 8)}
assert_model_is_valid(
prog,
inputs,
expected_output_shapes={block.outputs[0].name: (1, 2, 8, 10)},
)
def test_slicebyindex_mask_elimination(begin_mask, end_mask):
@mb.program(input_specs=[mb.TensorSpec(shape=(4, 4))])
def prog(x):
begin = [1, 1]
end = [1, 1]
for i in range(2):
if not begin_mask[i]:
begin[i] = 0
if not end_mask[i]:
end[i] = 4
r1 = mb.slice_by_index(x=x,
begin=begin,
end=end,
begin_mask=begin_mask,
end_mask=end_mask)
return mb.relu(x=r1)
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::noop_elimination")
assert get_op_types_in_program(prev_prog) == ["slice_by_index", "relu"]
assert get_op_types_in_program(prog) == ["relu"]
assert_model_is_valid(
prog,
{"x": (4, 4)},
expected_output_shapes={block.outputs[0].name: (4, 4)},
)
def test_gelu_tanh_approximation():
"""
Detect gelu tanh approx pattern, found in the TF bert model.
y = ( tanh((.0447)x^3 + x ) * (sqrt(2/pi)) + 1 ) * 0.5 * x
"""
@mb.program(input_specs=[mb.TensorSpec(shape=(3, 5, 6))])
def prog(x):
x1 = mb.pow(x=x, y=3)
x1 = mb.mul(x=0.044715, y=x1)
x1 = mb.add(x=x1, y=x)
x1 = mb.mul(x=x1, y=np.sqrt(2 / np.pi))
x1 = mb.tanh(x=x1)
x1 = mb.add(x=1, y=x1)
x1 = mb.mul(x=0.5, y=x1)
x1 = mb.mul(x=x, y=x1)
return x1
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::fuse_gelu_tanh_approximation")
assert get_op_types_in_program(prev_prog) == [
"pow",
"mul",
"add",
"mul",
"tanh",
"add",
"mul",
"mul",
]
assert get_op_types_in_program(prog) == ["gelu"]
assert_model_is_valid(
prog,
{"x": (3, 5, 6)},
expected_output_shapes={block.outputs[0].name: (3, 5, 6)},
)
def test_program_bgr(self):
"""
Input graph:
main(x: ImageType(color_layout="BGR", channel_first=True)) {
y1 = relu(x)
y2 = relu(x)
output = add(y1, y2)
} [output]
Output graph:
main(x: ImageType(channel_first=True)) {
y1 = relu(x)
y2 = relu(x)
output = add(y1, y2)
} [output]
"""
@mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 20, 20))])
def prog(x):
y1 = mb.relu(x=x)
y2 = mb.relu(x=x)
z = mb.add(x=y1, y=y2)
return z
prog.main_input_types = (ct.ImageType(name='x',
shape=[1, 3, 20, 20],
color_layout="BGR",
channel_first=True), )
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "mil_backend::insert_image_preprocessing_ops")
assert get_op_types_in_program(prev_prog) == ["relu", "relu", "add"]
assert get_op_types_in_program(prog) == ["relu", "relu", "add"]
def test_add_conv_transpose_output_shape():
"""
Given:
%1: (1, 5, 39, fp32) = conv_transpose(...) # no output_shape input.
Result:
%2: (3, i32) = const(val=[1,5,39])
%3: (1, 5, 39, fp32) = conv_transpose(..., output_shape=%2)
"""
N, C_in, C_out, D1 = 1, 3, 5, 20
@mb.program(input_specs=[mb.TensorSpec(shape=(N, C_in, D1))])
def prog(x):
weight = np.random.rand(C_out, C_in, D1).astype(np.float32)
return mb.conv_transpose(x=x, weight=weight)
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::add_conv_transpose_output_shape")
assert get_op_types_in_program(prev_prog) == ["conv_transpose"]
assert get_op_types_in_program(prog) == ["conv_transpose"]
prev_conv_transpose_op = prev_prog.find_ops(op_type="conv_transpose",
exactly_one=True)[0]
conv_transpose_op = prog.find_ops(op_type="conv_transpose",
exactly_one=True)[0]
assert np.all(conv_transpose_op.output_shape.val ==
prev_conv_transpose_op.outputs[0].shape)
def test_mixed_input_dtypes(self, op, x_dtype, y_dtype):
@mb.program(input_specs=[
mb.TensorSpec(shape=(10, 10), dtype=string_to_builtin(x_dtype)),
mb.TensorSpec(shape=(10, 10), dtype=string_to_builtin(y_dtype))
])
def prog(x, y):
x = getattr(mb, op)(x=x, y=y)
return x
assert get_op_types_in_program(prog) == [op]
_, _, block = apply_pass_and_basic_check(
prog, "mil_backend::homogenize_input_dtypes")
assert get_op_types_in_program(prog) == ["cast", op]
promoted_dtype = promote_types(string_to_builtin(x_dtype),
string_to_builtin(y_dtype))
# Asserting cast configuration
cast = block.find_ops(op_type="cast")[0]
assert cast.dtype.val == builtin_to_string(promoted_dtype)
assert len(cast.outputs) == 1
assert len(cast.outputs[0].child_ops) == 1
assert cast.outputs[0].child_ops[0].op_type == op
def test_adjust_cast(self):
"""
Input graph:
func main(int32 x) {
fp64 y = cast(x=x, dtype="fp64")
} -> (y)
becomes
func main(int32 x) {
fp32 y = cast(x=x, dtype="fp32")
} -> (y)
"""
@mb.program(
input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.int32)])
def prog(x):
y = mb.cast(x=x, dtype="fp64")
return y
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "mil_backend::adjust_io_to_supported_types")
assert get_op_types_in_program(prev_prog) == ['cast']
assert get_op_types_in_program(prog) == ['cast']
prev_cast = prev_prog.functions['main'].operations[1]
cast = prog.functions['main'].operations[2]
assert prev_cast.dtype.val == "fp64"
assert prev_cast.outputs[0].dtype == types.fp64
assert cast.dtype.val == "fp32"
assert cast.outputs[0].dtype == types.fp32
def test_onehot_matmul_to_gather_fusion(rank):
"""
Input:
%2 = one_hot(%1, on_value=1, off_value=0, axis=-1)
%3 = const() # rank 2
%4 = matmul(%2, %3)
Output:
%4 = gather(%3, %2, axis=0)
"""
rank4_shape = (10, 3, 6, 7)
input_shape = rank4_shape[-rank:]
vocab_size = 15
embedding_size = 12
@mb.program(input_specs=[mb.TensorSpec(shape=input_shape, dtype=types.int32)])
def prog(x):
x = mb.one_hot(
indices=x, on_value=1, off_value=0, axis=-1, one_hot_vector_size=vocab_size
)
x = mb.matmul(x=x, y=np.random.rand(vocab_size, embedding_size))
return x
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::fuse_onehot_matmul_to_gather"
)
assert get_op_types_in_program(prev_prog) == ["one_hot", "matmul"]
assert get_op_types_in_program(prog) == ["gather"]
assert_model_is_valid(
prog,
{"x": input_shape},
expected_output_shapes={block.outputs[0].name: input_shape + (embedding_size,)},
)
def test_pad_transposed_forked_conv(self):
@mb.program(input_specs=[mb.TensorSpec(shape=(1, 16, 20, 24))])
def prog(x):
pad = mb.pad(x=x, pad=[0, 0, 1, 1, 1, 1, 0, 0])
x = mb.transpose(x=pad, perm=[0, 3, 1, 2])
x = mb.conv(x=x,
weight=np.random.random([24, 24, 3, 3]),
pad_type="valid")
x = mb.transpose(x=x, perm=[0, 2, 3, 1])
y = mb.transpose(x=pad, perm=[0, 3, 1, 2])
y = mb.conv(x=y,
weight=np.random.random([24, 24, 3, 3]),
pad_type="valid")
y = mb.transpose(x=y, perm=[0, 2, 3, 1])
return x, y
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::pad_conv_connect")
self.assertEqual(get_op_types_in_program(prev_prog), [
"pad", "transpose", "conv", "transpose", "transpose", "conv",
"transpose"
])
self.assertEqual(get_op_types_in_program(prog), [
"transpose", "pad", "conv", "transpose", "transpose", "pad",
"conv", "transpose"
])
assert_model_is_valid(
prog,
{"x": (1, 16, 20, 24)},
expected_output_shapes={
block.outputs[0].name: (1, 16, 20, 24),
block.outputs[1].name: (1, 16, 20, 24)
},
)
def test_weight_decompression():
"""
This test is doing the following steps
(1) compress a model with two conv layers into a compressed model with two different constexpr ops
[Original model]:
weight_1 weight_2
| |
v v
input -> conv_1 -----> conv_2 ---> output
[Compressed model]:
weight_1_lut weight_2_affine
| |
v v
input -> conv_1 ------> conv_2 ---> output
, where weight_1_lut is a constexpr_lut_to_dense op and weight_2_affine is a constexpr_affine_dequantize op
(2) decompress the compressed model
[Decompressed model]:
weight_1_new weight_2_new
| |
v v
input -> conv_1 ------> conv_2 ---> output
, note that, weight_1_new is equivalent to weight_1_lut, and weight_2_new is equivalent to weight_2_affine
"""
model, inputs, torch_input_values, coreml_input_values = get_test_model_and_data(multi_layer=True)
torchmodel = torch.jit.trace(model, torch_input_values)
mlmodel = ct.convert(torchmodel, inputs=inputs, convert_to="mlprogram")
# we first compress the model
mlmodel = ct.compression_utils.palettize_weights(mlmodel, mode="kmeans", nbits=4, op_selector=lambda const_op: const_op.name == "conv_1_weight_to_fp16")
mlmodel = ct.compression_utils.affine_quantize_weights(mlmodel, mode="linear", op_selector=lambda const_op: const_op.name == "conv_2_weight_to_fp16")
expected_ops = ['constexpr_lut_to_dense', 'cast', 'conv', 'constexpr_affine_dequantize', 'conv', 'cast']
assert get_op_types_in_program(mlmodel._mil_program) == expected_ops
# decompress the model
decompressed_model = ct.compression_utils.decompress_weights(mlmodel)
assert get_op_types_in_program(decompressed_model._mil_program) == ['cast', 'conv', 'conv', 'cast']
if ct.utils._macos_version() < (13, 0):
return
# compared the numerical outputs
output_dict = mlmodel.predict(coreml_input_values)
de_output_dict = decompressed_model.predict(coreml_input_values)
for k, v in output_dict.items():
assert k in de_output_dict
np.testing.assert_allclose(v, de_output_dict[k])
def test_conv(self, rank, groups, has_bias, backend):
"""
Input graph:
input -----> conv -----> batch_norm ---> out
Output graph:
input -----> conv ----> out
"""
Cin, Cout = 10, 30
input_shape = (2, Cin, 20) if rank == 3 else (2, Cin, 20, 24)
@mb.program(input_specs=[mb.TensorSpec(shape=input_shape)])
def prog(x):
# conv layer
conv_weight = np.random.rand(Cout, Cin // groups, 2) if rank == 3 else np.random.rand(Cout, Cin // groups, 2, 3)
conv_bias = np.random.rand(Cout) if has_bias else None
x = mb.conv(
x=x,
weight=conv_weight,
bias=conv_bias,
groups=groups,
)
# batch_norm layer
gamma = np.random.rand(Cout)
beta = np.random.rand(Cout)
mean = np.random.rand(Cout)
variance = np.random.rand(Cout)
epsilon = 1e-2
x = mb.batch_norm(
x=x,
mean=mean,
variance=variance,
gamma=gamma,
beta=beta,
epsilon=epsilon,
)
return x
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::fuse_conv_batchnorm"
)
assert get_op_types_in_program(prev_prog) == ["conv", "batch_norm"]
assert get_op_types_in_program(prog) == ["conv"]
# validate graph pass
input_dict = {
"x": np.random.rand(*input_shape),
}
output_shape = (2, Cout, 19) if rank == 3 else (2, Cout, 19, 22)
assert_model_is_valid(
prog,
{"x": input_shape},
expected_output_shapes={block.outputs[0].name: output_shape},
backend=backend,
)
def test_conv_transpose(self, rank, groups, has_bias, scale_op, scale_type,
backend):
"""
Input graph:
input -----> conv_transpose -----> mul/real_div ---> out
Output graph:
input -----> conv_transpose ----> out
"""
Cin, Cout = 10, 30
input_shape = (2, Cin, 20) if rank == 3 else (2, Cin, 20, 24)
@mb.program(input_specs=[mb.TensorSpec(shape=input_shape)])
def prog(x):
# conv layer
conv_weight = np.random.rand(Cin, Cout // groups,
2) if rank == 3 else np.random.rand(
Cin, Cout // groups, 2, 3)
conv_bias = np.random.rand(Cout) if has_bias else None
x = mb.conv_transpose(
x=x,
weight=conv_weight,
bias=conv_bias,
groups=groups,
)
if scale_type == "scalar":
scale = np.array([2.3])
else:
scale = np.arange(Cout).astype(np.float32)
scale = np.reshape(scale, (Cout, 1) if rank == 3 else
(1, Cout, 1, 1))
# scale layer
if scale_op == "mul":
x = mb.mul(x=x, y=scale)
elif scale_op == "real_div":
x = mb.real_div(x=x, y=scale)
return x
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::fuse_conv_scale")
assert get_op_types_in_program(prev_prog) == [
"conv_transpose", scale_op
]
assert get_op_types_in_program(prog) == ["conv_transpose"]
# validate graph pass
output_shape = (2, Cout, 21) if rank == 3 else (2, Cout, 21, 26)
assert_model_is_valid(
prog,
{"x": input_shape},
expected_output_shapes={block.outputs[0].name: output_shape},
backend=backend,
)
def test_with_interleave(self):
"""
input1(1, 5, 3, 4) -----> stack(axis=2) -----> reshape(shape=(1, 10, 3, 4)) ---> out(1, 10, 3, 4)
^
|
input2(1, 5, 3, 4) ----------
Output graph:
input -----> concat ----> out
"""
@mb.program(input_specs=[mb.TensorSpec(shape=(1, 5, 3, 4)), mb.TensorSpec(shape=(1, 5, 3, 4))])
def prog(x1, x2):
x = mb.stack(values=[x1, x2], axis=2)
x = mb.reshape(x=x, shape=[1, 10, 3, 4])
return x
prev_prog, prev_block, block = apply_pass_and_basic_check(
prog, "common::replace_stack_reshape"
)
self.assertEqual(
get_op_types_in_program(prev_prog), ["stack", "reshape"]
)
self.assertEqual(get_op_types_in_program(prog), ["concat"])
inputs = {"x1": (1, 5, 3, 4), "x2": (1, 5, 3, 4)}
assert_model_is_valid(
prog,
inputs,
expected_output_shapes={block.outputs[0].name: (1, 10, 3, 4)},
)
concat_ops = [op for op in block.operations if op.op_type == 'concat']
concat_op = concat_ops[0]
assert concat_op.interleave.val == True
output_name = block.outputs[0].name
mlmodel = ct.convert(prog, source="milinternal", convert_to="neuralnetwork")
if not _IS_MACOS:
# Can not get predictions unless on macOS.
return
input_dict = dict()
for name, shape in inputs.items():
input_dict[name] = np.random.rand(*shape)
old_prediction = np.reshape(np.stack([input_dict["x1"], input_dict["x2"]], axis=2), newshape=[1, 10, 3, 4])
prediction = mlmodel.predict(input_dict, useCPUOnly=True)
np.testing.assert_allclose(old_prediction, prediction[output_name], atol=1e-04, rtol=1e-05)
