Exemple #1
0
    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)},
        )
Exemple #2
0
    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)
            },
        )
Exemple #3
0
    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),
            },
        )
Exemple #4
0
    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_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_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)},
        )
Exemple #8
0
    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)},
        )
Exemple #9
0
    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)},
        )
Exemple #10
0
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_nested_block(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, ))])
        def prog(x):
            def true_fn():
                # returns var with name x shadows input 'x'
                return mb.add(x=x, y=1, name='x')

            def false_fn():
                # two ops with name "x"
                return mb.add(x=x, y=-1, name='x')

            pred = mb.equal(x=mb.squeeze(x=x), y=1)
            return mb.cond(pred=pred, _true_fn=true_fn, _false_fn=false_fn)

        cond_op = prog.functions['main'].operations[-1]
        assert cond_op.blocks[0].outputs[0].name == 'x'
        assert cond_op.blocks[1].outputs[0].name == 'x'
        prev_prog, _, block = apply_pass_and_basic_check(
            prog, "common::dedup_op_and_var_names")
        cond_op = prog.functions['main'].operations[-1]
        assert cond_op.blocks[0].outputs[0].name == 'x_1'
        assert cond_op.blocks[1].outputs[0].name == 'x_2'

        assert_model_is_valid(
            prog,
            {"x": (1, )},
            expected_output_shapes={block.outputs[0].name: (1, )},
        )
    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_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)},
    )
Exemple #14
0
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,)},
    )
Exemple #15
0
    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)},
        )
Exemple #16
0
    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_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_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)},
        )
Exemple #19
0
def test_remove_vacuous_cond():
    @mb.program(input_specs=[
        mb.TensorSpec(shape=(1, ), dtype=types.bool),
        mb.TensorSpec(shape=(2, 3)),
    ])
    def prog(a, b):
        def then_branch():
            return mb.identity(x=b)

        def else_branch():
            return mb.identity(x=b)

        pred = mb.squeeze(x=a)
        return mb.cond(pred=pred, _true_fn=then_branch, _false_fn=else_branch)

    cond_op = prog.find_ops(op_type="cond", exactly_one=True)[0]
    original_cond_op_name = cond_op.name
    assert len(cond_op.blocks[0].operations) == 1
    assert len(cond_op.blocks[1].operations) == 1
    assert cond_op.blocks[0].operations[0].op_type == "identity"
    assert cond_op.blocks[1].operations[0].op_type == "identity"

    prev_prog = copy.deepcopy(prog)
    PASS_REGISTRY["tensorflow2::remove_vacuous_cond"](prog)
    assert_same_output_names(prev_prog, prog)

    cond_op = prog.find_ops(op_type="cond")
    assert len(cond_op) == 0
    identity_op = prog.find_ops(prefix=original_cond_op_name,
                                exactly_one=True)[0]
    assert identity_op.op_type == "identity"

    if validate_model:
        assert_model_is_valid(prog, {"a": (1, ), "b": (2, 3)})
Exemple #20
0
    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)},
        )
Exemple #21
0
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_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)
    def test_multiple(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 3, 4)), mb.TensorSpec(shape=(1, 2, 3, 4)), 
                                 mb.TensorSpec(shape=(1, 2, 3, 4)), mb.TensorSpec(shape=(1, 2, 3, 4))])
        def prog(x1, x2, x3, x4):
            a = mb.stack(values=[x1, x2], axis=1)
            a = mb.reshape(x=a, shape=[1, 4, 3, 4])

            b = mb.stack(values=[x3, x4], axis=1)
            b = mb.reshape(x=b, shape=[1, 4, 3, 4])

            c = mb.stack(values=[a, b], axis=2) 
            c = mb.reshape(x=c, shape=[1, 4, 6, 4])

            return c

        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", "stack", "reshape", "stack", "reshape"]
        )
        self.assertEqual(get_op_types_in_program(prog), ["concat", "concat", "concat"])

        inputs = {"x1": (1, 2, 3, 4), "x2": (1, 2, 3, 4), "x3": (1, 2, 3, 4), "x4": (1, 2, 3, 4)}
        assert_model_is_valid(
            prog,
            inputs,
            expected_output_shapes={block.outputs[0].name: (1, 4, 6, 4)},
        )

        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)

        branch_1 = np.reshape(np.stack([input_dict['x1'], input_dict['x2']], axis=1), newshape=[1, 4, 3, 4])
        branch_2 = np.reshape(np.stack([input_dict['x3'], input_dict['x4']], axis=1), newshape=[1, 4, 3, 4])
        old_prediction = np.reshape(np.stack([branch_1, branch_2], axis=2), newshape=[1, 4, 6, 4])

        prediction = mlmodel.predict(input_dict, useCPUOnly=True)

        np.testing.assert_allclose(old_prediction, prediction[output_name], atol=1e-04, rtol=1e-05)
Exemple #26
0
def test_const_elimination():
    @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
    def prog(x):
        a = np.random.rand(2, 4).astype(np.float32)
        double_a = mb.add(x=a, y=a)
        return mb.add(x=x, y=double_a)

    assert_op_count_match(prog, expect=2, op="const")
    prev_prog = copy.deepcopy(prog)
    PASS_REGISTRY["common::const_elimination"](prog)
    assert_same_output_names(prev_prog, prog)
    assert_op_count_match(prog, expect=3, op="const")

    if validate_model:
        assert_model_is_valid(prog, {"x": (2, 4)})
def test_keep_tile():
    @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
    def prog(x):
        r1 = mb.tile(x=x, reps=[2, 2])
        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) == ["tile", "relu"]
    assert get_op_types_in_program(prog) == ["tile", "relu"]
    assert_model_is_valid(
        prog,
        {"x": (2, 4)},
        expected_output_shapes={block.outputs[0].name: (4, 8)},
    )
def test_linear_elimination():
    @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
    def prog(x):
        r1 = mb.linear_activation(x=x, alpha=1.0, beta=0.0)
        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) == ["linear_activation", "relu"]
    assert get_op_types_in_program(prog) == ["relu"]
    assert_model_is_valid(
        prog,
        {"x": (2, 4)},
        expected_output_shapes={block.outputs[0].name: (2, 4)},
    )
def test_crop_elimination():
    @mb.program(input_specs=[mb.TensorSpec(shape=(3, 2, 4))])
    def prog(x):
        r1 = mb.crop(x=x, crop_height=[0, 0], crop_width=[0, 0])
        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) == ["crop", "relu"]
    assert get_op_types_in_program(prog) == ["relu"]
    assert_model_is_valid(
        prog,
        {"x": (3, 2, 4)},
        expected_output_shapes={block.outputs[0].name: (3, 2, 4)},
    )
def test_upsample_bilinear_elimination():
    @mb.program(input_specs=[mb.TensorSpec(shape=(3, 2, 4))])
    def prog(x):
        r1 = mb.upsample_bilinear(x=x)
        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) == ["upsample_bilinear", "relu"]
    assert get_op_types_in_program(prog) == ["relu"]
    assert_model_is_valid(
        prog,
        {"x": (3, 2, 4)},
        expected_output_shapes={block.outputs[0].name: (3, 2, 4)},
    )