lesson 4 - assignment 疑问

[hbsun2113]

对于第一个题目：

我想使用使用如下方法来解决，把pytorch_model包装进函数Pytorch2IRModule，然后使用emit_te来生成对应的IRModule：

def Pytorch2IRModule(x):
  model_torch = pytorch_model()
  return model_torch(x)

def create_model_via_emit_te():
    bb = relax.BlockBuilder()
    x = relax.Var("x", input_shape, relax.DynTensorType(batch_size, "float32"))

    conv2d_weight = relax.const(weight_map["conv2d_weight"], "float32")
    conv2d_bias = relax.const(weight_map["conv2d_bias"].reshape(1, 32, 1, 1), "float32")
    linear0_weight = relax.const(weight_map["linear0_weight"], "float32")
    linear0_bias = relax.const(weight_map["linear0_bias"].reshape(1, 100), "float32")
    linear1_weight = relax.const(weight_map["linear1_weight"], "float32")
    linear1_bias = relax.const(weight_map["linear1_bias"].reshape(1, 10), "float32")

    with bb.function("main", [x]):
        with bb.dataflow():
           lv0 = bb.emit_te(Pytorch2IRModule, x)
           gv = bb.emit_output(lv0)
        bb.emit_func_output(gv)

    return bb.get()

但是报错如下：

请问：

1. 解题思路是否正确？如果不对，不能包装一个Model的原因是什么？
2. 上述报错信息的原因是什么？

[jinhongyii]

1. 不正确。emit_te的第一个参数需要是一个返回tensor expression的函数，不能是pytorch模型
2. emit_te会调用传入的第一个参数，并把之后的参数作为输入传递。因此在这里实际调用了 Pytorch2IRModule(x)，而其中使用了model_torch(x)，model_torch是一个torch模型，而x是relax var，不是torch.tensor，无法传入torch模型中。

[hbsun2113]

谢谢回复！

1. 针对第一个问题，Pytorch2IRModule返回的是tensor expression(e.g. torch.Tensor)，所以能满足emit_te的要求？
2. 针对第二个问题，我们是否有办法把relax var转换成torch.tensor? 如果不能，原因是什么？

[jinhongyii]

1. 你对于tensor expression的理解可能出了差错，torch.Tensor不是tensor expression。关于tensor expression请看这里的notes：https://mlc.ai/chapter_tensor_program/case_study.html#generate-tensorir-code-using-tensor-expression。
2. relax var是relax IR里的变量，是抽象意义上的，而torch.Tensor是torch里的张量，它包含了数据本身。这两者完全没有任何联系。

[hbsun2113]

是的，我之前对tensor expression的理解不对。
另一个问题：

1. Tensor Expression和Primitive Tensor Function是什么关系？我理解Tensor Expression用于表达运算关系，而Primitive Tensor Function是例如 linear, add, relu, and softmax。所以我理解Tensor Expression是Primitive Tensor Function的实现？
2. 如果如上所言，那为什么说Tensor Expression可以用于简化Primitive Tensor Function？
   

[jinhongyii]

1. tensor expression是生成primitive tensor function的一种方式
2. 本来我们需要手写TVMScript来得到primitive tensor function，这很费劲。而使用tensor expression可以通过简单的几行代码快速生成想要的primitive tensor function。

[hbsun2113]

Hi @jinhongyii @Hzfengsy , 请教下emit会对结果进行什么转换？在完成assignment中作业1时我的code如下，可以看到myTorchModel::out_torch和output_from_pytorch输出已经完全相同了，但是经过emit后输出被改变。

@tvm.register_func("env.myTorchModel", override=True)
def myTorchModel(
    x: tvm.nd.NDArray,
    conv2d_weight: tvm.nd.NDArray,
    conv2d_bias: tvm.nd.NDArray,
    linear0_weight: tvm.nd.NDArray,
    linear0_bias: tvm.nd.NDArray,
    linear1_weight: tvm.nd.NDArray,
    linear1_bias: tvm.nd.NDArray,
    out: tvm.nd.NDArray
):
    x_torch = torch.from_dlpack(x)
    conv2d_weight_torch = torch.from_dlpack(conv2d_weight)
    conv2d_bias_torch = torch.from_dlpack(conv2d_bias)
    linear0_weight_torch = torch.from_dlpack(linear0_weight)
    linear0_bias_torch = torch.from_dlpack(linear0_bias)
    linear1_weight_torch = torch.from_dlpack(linear1_weight)
    linear1_bias_torch = torch.from_dlpack(linear1_bias)
    out_torch = torch.from_dlpack(out)
    
    layer0 = torch.nn.functional.conv2d(x_torch, conv2d_weight_torch, conv2d_bias_torch)
    layer1 = torch.nn.functional.relu(layer0)
    layer2 = torch.nn.functional.max_pool2d(layer1, kernel_size=(2, 2))
    layer3 = torch.flatten(layer2, start_dim=1)
    layer4 = torch.nn.functional.linear(layer3, linear0_weight_torch, linear0_bias_torch)
    layer5 = torch.nn.functional.relu(layer4)
    layer6 = torch.nn.functional.linear(layer5, linear1_weight_torch, linear1_bias_torch)
    layer7 = torch.nn.functional.softmax(layer6, dim=1)

    out_torch = layer7
    print("out_torch", out_torch.numpy().shape, out_torch)

# function to create IR with mixed model
def create_model_via_emit_te():
    bb = relax.BlockBuilder()

    x = relax.Var("x", input_shape, relax.DynTensorType(4, "float32"))

    conv2d_weight = relax.const(weight_map["conv2d_weight"].reshape(32, 1, 3, 3), "float32")
    conv2d_bias = relax.const(weight_map["conv2d_bias"].reshape(32, ), "float32")
    linear0_weight = relax.const(weight_map["linear0_weight"], "float32")
    linear0_bias = relax.const(weight_map["linear0_bias"].reshape(100), "float32")
    linear1_weight = relax.const(weight_map["linear1_weight"], "float32")
    linear1_bias = relax.const(weight_map["linear1_bias"].reshape(10), "float32")
    
    with bb.function("main", [x]):
        with bb.dataflow():
            lv0 = bb.emit(relax.op.call_tir(relax.extern("env.myTorchModel"), (x, conv2d_weight, conv2d_bias, linear0_weight, linear0_bias, linear1_weight, linear1_bias), (4, 10), dtype="float32"))
            gv = bb.emit_output(lv0)
        bb.emit_func_output(gv)

    return bb.get()


def build_mod(mod):
    exec = relax.vm.build(mod, "llvm")
    dev = tvm.cpu()
    vm = relax.VirtualMachine(exec, dev)
    return vm


def check_equivalence(mod, torch_model, test_loader):
    torch_model.eval()
    with torch.no_grad():
        rt_mod = build_mod(mod)
        for data, label in test_loader:
            data, label = data.cpu(), label.cpu()
            output_from_pytorch = torch_model(data).numpy()
            output_from_relax = rt_mod["main"](tvm.nd.array(data, tvm.cpu())).numpy()
            print("output_from_pytorch", output_from_pytorch.shape, output_from_pytorch)
            print("output_from_relax", output_from_relax.shape, output_from_relax)
            tvm.testing.assert_allclose(output_from_pytorch, output_from_relax, rtol=1e-4)


test_data = torchvision.datasets.FashionMNIST(
    "./data",
    download=True,
    train=False,
    transform=transforms.Compose([transforms.ToTensor()])
)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False)

mod = create_model_via_emit_te()
torch_model = pytorch_model()

check_equivalence(mod, torch_model, test_loader)
IPython.display.Code(mod.script(), language="python")

输出：红色和黄色相同，蓝色不同

请问：

1. 发生变化的原因是什么？
2. 我该如何改正这个错误？

[hbsun2113]

attach the script:

[Hzfengsy]

out_torch = torch.from_dlpack(out)
out_torch = layer7

和课程中的

out_torch = torch.from_dlpack(out)
torch.add(out_torch, b_torch, out=out_torch)

不等价（可以自己思考一下不等价的原因）

[hbsun2113]

1. 使用torch.add(layer7, 0, out=out_torch)可以通过测试。
2. 对于您说的不等价，我猜测out_torch = layer7使得out_torch 更改了指向位置，即不再指向call_tir为其开辟的空间，而torch.add(out_torch, b_torch, out=out_torch)是仍向call_tir所开辟空间的位置填写结果。不知道是否理解正确？
3. 如果把torch.add(layer7, 0, out=out_torch)替换为out_torch = layer7 + 0则test failed，也是因为更改了指向位置，对吗？

[Hzfengsy]

你的理解基本上是对的：

1. out_torch = torch.from_dlpack(out)生成的out_torch是一个内存指针，指向了TVM生成的空间
2. torch.add(layer7, 0, out=out_torch)是指将计算结果写入out_torch生成的空间
3. out_torch = layer7 + 0是指生成一块含有计算空间的地址，并且将地址返回到out_torch。值得注意的是，虽然out_torch的结果是正确的，但是并没有写回TVM指定的内存空间