Function test_mlp_blockbuilder

tests/python/relax/test_transform_gradient_numeric.py:86–138

(target, dev)

@tvm.testing.parametrize_targets("llvm")
def test_mlp_blockbuilder(target, dev):
    layers, in_size, out_size, hidden_size, batch_size = 3, 5, 5, 5, 4

    input_list = [relax.Var("x", R.Tensor((batch_size, in_size), "float32"))]
    w_list = (
        [relax.Var("w_0", R.Tensor((in_size, hidden_size), "float32"))]
        + [
            relax.Var("w_" + str(i + 1), R.Tensor((hidden_size, hidden_size), "float32"))
            for i in range(layers - 2)
        ]
        + [relax.Var("w_" + str(layers - 1), R.Tensor((hidden_size, out_size), "float32"))]
    )
    b_list = [
        relax.Var("b_" + str(i), R.Tensor((hidden_size,), "float32")) for i in range(layers - 1)
    ] + [relax.Var("b_" + str(layers - 1), R.Tensor((out_size,), "float32"))]
    label_list = [relax.Var("y", R.Tensor((batch_size,), "int64"))]
    args_list = input_list + w_list + b_list + label_list

    bb = relax.BlockBuilder()
    with bb.function("MLP", args_list):
        with bb.dataflow():
            current = input_list[0]
            for i in range(layers):
                lv0 = bb.emit(R.matmul(current, w_list[i]))
                lv1 = bb.emit(R.add(lv0, b_list[i]))
                current = bb.emit(R.nn.relu(lv1) if i < layers - 1 else lv1)
            logits = R.nn.log_softmax(current)
            loss = bb.emit(R.nn.nll_loss(logits, label_list[0]))
            gv0 = bb.emit_output(loss)
        bb.emit_func_output(gv0)

    Before = bb.get()
    After = relax.transform.Gradient("MLP", w_list + b_list)(Before)
    # Check numerical gradients equal
    args = []
    for arg in After["MLP_adjoint"].params:
        shape = [int(l) for l in arg.struct_info.shape]
        if arg.struct_info.dtype == "int64":
            args.append(
                tvm.runtime.tensor(np.random.randint(0, out_size, size=shape).astype(np.int64))
            )
        else:  # float32
            args.append(rand("float32", *shape))

    vm_before = _legalize_and_build(Before, target, dev)
    vm_after = _legalize_and_build(After, target, dev)
    _, grad = vm_after["MLP_adjoint"](*args)

    def func(*inputs):
        loss = vm_before["MLP"](args[0], *[tvm.runtime.tensor(i) for i in inputs], args[-1])
        return loss.numpy()

    check_numerical_grads(func, [i.numpy() for i in args[1:-1]], [i.numpy() for i in grad])


@tvm.testing.parametrize_targets("llvm")