A error occurred when I run a code snippet which is a small self-attenion implementation. I also disable ‘AlterOpLayout’ pass,it doesn’t work. Is my code wrong?How to modify it. Here is my source code
from tvm import relay
from tvm.relay import testing
import tvm
from tvm import te
from tvm.contrib import graph_executor
import numpy as np
def dense(data, weight=None, units=None, **kwargs):
name = kwargs.get("name")
kwargs.pop("name")
if not weight:
weight = relay.var(name + "_weight")
ans = relay.nn.dense(data, weight, units, **kwargs)
return ans
def batch_matmul(data, weight=None,**kwargs):
return relay.nn.batch_matmul(data,weight)
def bert(seq_lenth=64,hidden_dims=1024,heads=64):
data = relay.var("data", shape=(seq_lenth,hidden_dims), dtype='float32')
body = relay.nn.relu(data=data)
# attention
d_k = hidden_dims // heads
q = dense(data,units=hidden_dims,name="q")
k = dense(data,units=hidden_dims,name="k")
v = dense(data,units=hidden_dims,name="v")
q = relay.reshape(q,(-1,heads,d_k))
k = relay.reshape(k,(-1,heads,d_k))
q = relay.reshape(q,(-1,heads,d_k))
q = relay.transpose(q,(1,0,2))
k = relay.transpose(k,(1,0,2))
v = relay.transpose(v,(1,2,0))
logits = batch_matmul(q,k,name='logits')
output = batch_matmul(logits,v,name='output')
output = relay.transpose(output,(1,0,2))
output = relay.reshape(output,(seq_lenth,hidden_dims))
net = dense(output,units=hidden_dims,name='net')
return relay.Function(relay.analysis.free_vars(net), net)
def get_workload():
net = bert()
mod = tvm.IRModule.from_expr(net)
mod = relay.transform.InferType()(mod)
shape_dict = {v.name_hint: v.checked_type for v in mod["main"].params}
np.random.seed(0)
params = {}
for k, v in shape_dict.items():
if k == "data":
continue
init_value = np.random.uniform(-1, 1, v.concrete_shape).astype(v.dtype)
params[k] = tvm.nd.array(init_value, device=tvm.cpu(0))
return mod,params
if __name__=='__main__':
mod,params = get_workload()
data_shape = (64,1024)
opt_level = 4
target = tvm.target.cuda()
with tvm.transform.PassContext(opt_level=opt_level,disabled_pass={"AlterOpLayout"}):
lib = relay.build(mod, target, params=params)
dev = tvm.cuda()
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
# create module
module = graph_executor.GraphModule(lib["default"](dev))
# set input and parameters
module.set_input("data", data)
# run
print(module.benchmark(dev, number=1, repeat=600))
And error message
Check failed: (!axes.defined() || static_cast<int>(axes.size()) == ndim) is false: Dimension mismatch: axes has 3 elements, but data.ndim = 2
And I think it failed at
output = batch_matmul(logits,v,name='output')
Thanks a lot!