TVM compilation failed on SimpleViT
yangxin0926 opened this issue · 0 comments
yangxin0926 commented
I use this repo to export a SimpleViT onnx model. And I tried to compile the onnx model with relay but I got the following error.
Check failed: src_idx < ishape.size() (3 vs. 3) :
I thought this was raised by multihead for after changing heads to 1 the error does not occur.
But after I changed the heads to 1, I've got another error.
Check failed: (reporter->AssertEQ(xk, yk)) is false: BatchDot: shapes of x and y is inconsistent, x shape=[1, 64, 1], y shape=[1, 64, 1024]
I think this error is raised by node /transformer/layers.0.1/net/net.0/ReduceMean
because it's the last node that relay processed.
I don't know what to solve these problems. Here is my code generating onnx model.
from vit_pytorch import SimpleViT
image_size = 256
patch_size = 32
num_classes = 1000
dim = 1024
depth = 6
heads = 1
mlp_dim = 2048
model = SimpleViT(
image_size = image_size,
patch_size = patch_size,
num_classes = num_classes,
dim = dim,
depth = depth,
heads = heads,
mlp_dim = mlp_dim
)
model.eval()
dummy_input = torch.randn(1, 3, image_size, image_size, requires_grad=True)
print(model)
print(model(dummy_input))
torch.onnx.export(
model,
dummy_input,
"ViT.onnx",
export_params = True,
opset_version = 11,
input_names = ['imageInput'],
output_names = ['predOutput'],
dynamic_axes = {'imageInput' : {0 : 'batch_size'}, 'predOutput' : {0 : 'batch_size'}})
And here is my code to compile the onnx model.
import onnx
import tvm
from tvm import te
import tvm.relay as relay
model = onnx.load("ViT.onnx")
shape_dict = {"imageInput" : [1, 3, 256, 256]}
mod, params = relay.frontend.from_onnx(model, shape_dict)