I’ve defined a conv2d operation using TVM script, which is equivalent to pytorch’s conv2d operation. But I am using numpy.test.allclose () to verify that the tensors I defined with tvm are different from the ones I ran with pytorch。Here is my code：
import numpy as np
import tvm
from tvm.script import tir as T
import torch
N, CI, H, W, CO, K = 1, 1, 8, 8, 2, 3
OUT_H, OUT_W = H - K + 1, W - K + 1
data = np.arange(N * CI * H * W).reshape(N, CI, H, W).astype(np.int64)
weight = np.arange(CO * CI * K * K).reshape(CO, CI, K, K).astype(np.int64)
data_torch = torch.Tensor(data)
weight_torch = torch.Tensor(weight)
print(data_torch)
print(weight_torch)
conv_torch = torch.nn.functional.conv2d(data_torch, weight_torch)
@tvm.script.ir_module
class MyConv:
@T.prim_func
def conv(d: T.Buffer[(1, 1, 8, 8), "int64"], w: T.Buffer[(2, 1, 3, 3), "int64"],
out: T.Buffer[(1, 2, 6, 6), "int64"]):
T.func_attr({"global_symbol": "conv", "tir.noalias": True})
for batch, co, ci, i, j, m, n in T.grid(1, 2, 1, 6, 6, 3, 3):
with T.block("out"):
v_batch = T.axis.spatial(1, batch)
v_co = T.axis.spatial(2, co)
v_ci = T.axis.spatial(1, ci)
vi = T.axis.spatial(6, i)
vj = T.axis.spatial(6, j)
vm = T.axis.spatial(3, m)
vn = T.axis.spatial(3, n)
out[v_batch, co, vi, vj] = out[v_batch, v_co, vi, vj] + \
d[batch, v_ci, vi + vm, vj + vn] * w[v_co, v_ci, vm, vn]
rt_lib = tvm.build(MyConv, target="llvm")
data_tvm = tvm.nd.array(data)
weight_tvm = tvm.nd.array(weight)
conv_tvm = tvm.nd.array(np.empty((N, CO, OUT_H, OUT_W), dtype=np.int64))
rt_lib["conv"](data_tvm, weight_tvm, conv_tvm)
print(conv_tvm.numpy())
np.testing.assert_allclose(conv_tvm.numpy(), conv_torch.numpy(), rtol=1e-5)
I find it fine to run assert monitoring directly if I don’t print the results initially。But if you print the result, the test goes wrong and the result isn’t correct, and if I print the pytorch result as well, suddenly the result is correct and consistent。But it still didn’t pass numpy’s test, which was a very strange phenomenon