def before(dim):
X = relay.var("X", shape=(1, dim))
W = relay.var("W", shape=(3 * dim, dim))
matmul = relay.nn.dense(X, W)
splitted = relay.split(matmul, indices_or_sections=3, axis=1)
out = relay.sigmoid(splitted[0]) + relay.tanh(splitted[1]) * relay.exp(splitted[2])
return relay.Function([X, W], out)
def get_net(batch_size, random_len=100, oshape=(3, 64, 64), ngf=128, code=None, dtype="float32"):
"""get net of dcgan generator"""
assert oshape[-1] == 64, "Only support 64x64 image"
assert oshape[-2] == 64, "Only support 64x64 image"
code = relay.var("data", dtype=dtype, shape=(batch_size, random_len)) if code is None else code
dense_weight = relay.var("dense_weight")
dense = relay.nn.dense(code, weight=dense_weight, units=4*4*ngf*8)
relu = relay.nn.relu(dense)
# 4 x 4
reshape = relay.reshape(relu, newshape=(-1, ngf * 8, 4, 4))
# 8 x 8
dc8 = deconv2d_bn_relu(
reshape, ishape=(ngf * 8, 4, 4), oshape=(ngf * 4, 8, 8), kshape=(4, 4), prefix="g2")
# 16x16
dc16 = deconv2d_bn_relu(
dc8, ishape=(ngf * 4, 8, 8), oshape=(ngf * 2, 16, 16), kshape=(4, 4), prefix="g3")
# 32x32
dc32 = deconv2d_bn_relu(
dc16, ishape=(ngf * 2, 16, 16), oshape=(ngf, 32, 32), kshape=(4, 4), prefix="g4")
# 64x64
dc64 = deconv2d(
dc32, ishape=(ngf, 32, 32), oshape=oshape[-3:], kshape=(4, 4), name="g5_deconv")
tanh = relay.tanh(dc64)
args = relay.ir_pass.free_vars(tanh)
return relay.Function(args, tanh)
def get_conv1d_bias(x_shape=(1, 3, 224), k_shape=(10, 3, 3), activation=None, dtype="float32"):
conv, dic, param_lst = get_conv1d(x_shape=x_shape, k_shape=k_shape, dtype=dtype)
bias = relay.var("bias", shape=(k_shape[0],), dtype=dtype)
out = relay.nn.bias_add(conv, bias)
dic["bias"] = (k_shape[0],)
param_lst += ["bias"]
if activation == "relu":
return relay.nn.relu(out), dic, param_lst
elif activation == "tanh":
return relay.tanh(out), dic, param_lst
elif activation == "sigmoid":
return relay.sigmoid(out), dic, param_lst
else:
return out, dic, param_lst
def _get_model(shape, input_zp, input_sc, output_zp, output_sc, dtype):
a = relay.var("a", shape=shape, dtype=dtype)
dequantize = relay.qnn.op.dequantize(
a,
input_scale=relay.const(input_sc, "float32"),
input_zero_point=relay.const(input_zp, "int32"),
)
tanh = relay.tanh(dequantize)
model = relay.qnn.op.quantize(
tanh,
output_scale=relay.const(output_sc, "float32"),
output_zero_point=relay.const(output_zp, "int32"),
out_dtype=dtype,
)
return model
def create_graph():
data = relay.var('data', shape=(10, 10))
bn_gamma = relay.var("bn_gamma")
bn_beta = relay.var("bn_beta")
bn_mmean = relay.var("bn_mean")
bn_mvar = relay.var("bn_var")
x = relay.nn.batch_norm(data, bn_gamma, bn_beta, bn_mmean, bn_mvar)
out_1 = relay.nn.relu(x[0])
bn_out_1 = x[1]
out_2 = relay.tanh(bn_out_1)
out_3 = relay.log(bn_out_1)
out = relay.Tuple([out_1, out_2, out_3])
func = relay.Function([data, bn_gamma, bn_beta, bn_mmean, bn_mvar],
out)
return func
def expected(dim):
p0 = relay.var("p0", shape=(1, dim))
p1 = relay.var("p1", shape=(3 * dim, dim))
matmul = relay.nn.dense(p0, p1)
f0 = relay.Function([p0, p1], matmul)
p01 = relay.var("p01", shape=(1, 3 * dim))
splitted = relay.split(p01, indices_or_sections=3, axis=1)
out = relay.sigmoid(splitted[0]) + relay.tanh(splitted[1]) * relay.exp(splitted[2])
f1 = relay.Function([p01], out)
X = relay.var("X", shape=(1, dim))
W = relay.var("W", shape=(3 * dim, dim))
y = relay.Call(f0, [X, W])
z = relay.Call(f1, [y])
return relay.Function([X, W], z)
def expected(dim):
p0 = relay.var("p0", shape=(1, dim))
p1 = relay.var("p1", shape=(3 * dim, dim))
matmul = relay.nn.dense(p0, p1)
f0 = relay.Function([p0, p1], matmul)
p01 = relay.var("p01", shape=(1, 3 * dim))
splitted = relay.split(p01, indices_or_sections=3, axis=1)
out = relay.sigmoid(splitted[0]) + relay.tanh(splitted[1]) * relay.exp(splitted[2])
f1 = relay.Function([p01], out)
X = relay.var("X", shape=(1, dim))
W = relay.var("W", shape=(3 * dim, dim))
y = relay.Call(f0, [X, W])
z = relay.Call(f1, [y])
return relay.Function([X, W], z)
def expected():
data = relay.var('data', shape=(10, 10))
cb_1 = compiler_begin(data, "test")
O_1 = relay.abs(cb_1)
ce_2 = compiler_end(O_1, "test")
O_2 = relay.nn.relu(O_1)
ce_3 = compiler_end(O_2, "test")
X = relay.tanh(ce_2)
cb_3 = compiler_begin(ce_3, "test")
cb_4 = compiler_begin(X, "test")
O_3 = relay.add(cb_3, cb_4)
ce_4 = compiler_end(O_3, "test")
func = relay.Function([data], ce_4)
return func
def create_external_func1(mod_, compiler_name, symbol_name):
x_int = relay.var("x_int", shape=(10, 10))
p0 = relay.nn.relu(x_int)
q0 = relay.tanh(x_int)
# reshapes
p0_reshaped = relay.reshape(p0, newshape=100)
q0_reshaped = relay.reshape(q0, newshape=100)
ofms = relay.concatenate((p0_reshaped, q0_reshaped), 0)
f1 = relay.Function([x_int], ofms)
f1 = set_func_attr(f1, compiler_name, symbol_name)
glb_f1 = relay.GlobalVar(symbol_name)
mod_[glb_f1] = f1
mod_ = relay.transform.InferType()(mod_)
return glb_f1, mod_
def test_region_set_creator_diamond():
data = relay.var('data', shape=(10, 10))
cb_1 = compiler_begin(data, 'test_target')
O_1 = relay.abs(cb_1)
ce_1 = compiler_end(O_1, 'test_target')
ce_2 = compiler_end(O_1, 'test_target')
cb_2 = compiler_begin(ce_1, 'test_target')
O_2 = relay.nn.relu(cb_2)
ce_3 = compiler_end(O_2, 'test_target')
cb_d = compiler_begin(ce_2, "default")
X = relay.tanh(cb_d)
ce_d = compiler_end(X, 'default')
cb_3 = compiler_begin(ce_3, 'test_target')
cb_4 = compiler_begin(ce_d, 'test_target')
O_3 = relay.add(cb_3, cb_4)
ce_4 = compiler_end(O_3, 'test_target')
diamond = relay.Function([data], ce_4)
region_set = relay.analysis.AnnotatedRegionSet(diamond,
relay.op.get("annotation.compiler_begin"),
relay.op.get("annotation.compiler_end"))
assert len(region_set) == 4
check_region(
region_set,
[cb_1],
[cb_1, O_1, ce_1, ce_2],
[ce_1, ce_2],
)
check_region(
region_set,
[cb_2],
[cb_2, O_2, ce_3],
[ce_3],
)
check_region(
region_set,
[cb_d],
[cb_d, X, ce_d],
[ce_d],
)
check_region(
region_set,
[cb_3, cb_4],
[cb_3, cb_4, O_3, ce_4],
[ce_4],
)
def conv2d(x, node: Conv, params):
wname = f'{node.name}.weight'
bname = f'{node.name}.bias'
weight = tvm.nd.array(node.weight)
bias = tvm.nd.array(node.bias)
params[wname] = weight
params[bname] = bias
x = relay.nn.conv2d(x, weight=relay.var(wname, shape=weight.shape), strides=node.stride, padding=node.padding, channels=node.out_channels, kernel_size=node.kernel, groups=node.groups)
x = relay.nn.bias_add(x, relay.var(bname, shape=bias.shape), axis=1)
if node.act == "relu":
x = relay.nn.relu(x)
elif node.act == "tanh":
x = relay.tanh(x)
elif node.act == "identity":
x = x
return x
def get_net(batch_size,
random_len=100,
oshape=(3, 64, 64),
ngf=128,
code=None,
dtype="float32"):
"""get net of dcgan generator"""
assert oshape[-1] == 64, "Only support 64x64 image"
assert oshape[-2] == 64, "Only support 64x64 image"
code = relay.var("data", dtype=dtype,
shape=(batch_size, random_len)) if code is None else code
dense_weight = relay.var("dense_weight")
dense = relay.nn.dense(code, weight=dense_weight, units=4 * 4 * ngf * 8)
relu = relay.nn.relu(dense)
# 4 x 4
reshape = relay.reshape(relu, newshape=(-1, ngf * 8, 4, 4))
# 8 x 8
dc8 = deconv2d_bn_relu(reshape,
ishape=(ngf * 8, 4, 4),
oshape=(ngf * 4, 8, 8),
kshape=(4, 4),
prefix="g2")
# 16x16
dc16 = deconv2d_bn_relu(dc8,
ishape=(ngf * 4, 8, 8),
oshape=(ngf * 2, 16, 16),
kshape=(4, 4),
prefix="g3")
# 32x32
dc32 = deconv2d_bn_relu(dc16,
ishape=(ngf * 2, 16, 16),
oshape=(ngf, 32, 32),
kshape=(4, 4),
prefix="g4")
# 64x64
dc64 = deconv2d(dc32,
ishape=(ngf, 32, 32),
oshape=oshape[-3:],
kshape=(4, 4),
name="g5_deconv")
tanh = relay.tanh(dc64)
args = relay.analysis.free_vars(tanh)
return relay.Function(args, tanh)
def create_graph():
data = relay.var('data', shape=(10, 10))
cb_1 = compiler_begin(data, 'test_target')
O_1 = relay.abs(cb_1)
ce_2 = compiler_end(O_1, 'test_target')
O_2 = relay.nn.relu(O_1)
ce_3 = compiler_end(O_2, 'test_target')
X = relay.tanh(ce_2)
cb_3 = compiler_begin(ce_3, 'test_target')
cb_4 = compiler_begin(X, 'test_target')
O_3 = relay.add(cb_3, cb_4)
ce_4 = compiler_end(O_3, 'test_target')
func = relay.Function([data], ce_4)
return func
def expected(dim):
p0 = relay.var("p0", shape=(1, dim))
p1 = relay.var("p1", shape=(3 * dim, dim))
matmul = relay.nn.dense(p0, p1)
f0 = relay.Function([p0, p1], matmul)
f0 = f0.set_attribute("Primitive", tvm.tir.IntImm("int32", 1))
p01 = relay.var("p01", shape=(1, 3 * dim))
splitted = relay.split(p01, indices_or_sections=3, axis=1)
out = relay.sigmoid(splitted[0]) + relay.tanh(splitted[1]) * relay.exp(splitted[2])
f1 = relay.Function([p01], out)
f1 = f1.set_attribute("Primitive", tvm.tir.IntImm("int32", 1))
X = relay.var("X", shape=(1, dim))
W = relay.var("W", shape=(3 * dim, dim))
y = relay.Call(f0, [X, W])
z = relay.Call(f1, [y])
return relay.Function([X, W], z)
def expected():
mod = tvm.IRModule()
# function 1
f1_cb1 = relay.var('test_target_1_i0', shape=(10, 10))
f1_O_1 = relay.abs(f1_cb1)
f1_O_2 = relay.nn.relu(f1_O_1)
f1_out = relay.Tuple((f1_O_2, f1_O_1))
func1 = relay.Function([f1_cb1], f1_out)
func1 = func1.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func1 = func1.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func1 = func1.with_attr("Compiler", tvm.tir.StringImm("test_target"))
func1 = func1.with_attr("global_symbol",
container.String("test_target_1"))
gv1 = relay.GlobalVar("test_target_1")
mod[gv1] = func1
# function 0
f2_cb3 = relay.var('test_target_0_i0', shape=(10, 10))
f2_cb4 = relay.var('test_target_0_i1', shape=(10, 10))
f2_O_3 = relay.add(f2_cb3, f2_cb4)
func0 = relay.Function([f2_cb3, f2_cb4], f2_O_3)
func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Compiler", tvm.tir.StringImm("test_target"))
func0 = func0.with_attr("global_symbol",
container.String("test_target_0"))
gv0 = relay.GlobalVar("test_target_0")
mod[gv0] = func0
# body
data = relay.var('data', shape=(10, 10))
tuple_out = gv1(data)
ce_2 = relay.TupleGetItem(tuple_out, 1)
ce_3 = relay.TupleGetItem(tuple_out, 0)
X = relay.tanh(ce_2)
ce_4 = gv0(ce_3, X)
func = relay.Function([data], ce_4)
mod["main"] = func
return mod
def diamond_graph_fanouts():
data = relay.var('data', shape=(10, 10))
cb_1 = compiler_begin(data, "test")
O_1 = relay.abs(cb_1)
ce_1 = compiler_end(O_1, "test")
ce_2 = compiler_end(O_1, "test")
cb_2 = compiler_begin(ce_1, "test")
O_2 = relay.nn.relu(cb_2)
ce_3 = compiler_end(O_2, "test")
X = relay.tanh(ce_2)
cb_3 = compiler_begin(ce_3, "test")
cb_4 = compiler_begin(X, "test")
O_3 = relay.add(cb_3, cb_4)
ce_4 = compiler_end(O_3, "test")
diamond = relay.Function([data], ce_4)
return diamond
def test_region_set_creator_merged():
data = relay.var("data", shape=(10, 10))
cb_1 = compiler_begin(data, "test_target")
O_1 = relay.abs(cb_1)
ce_2 = compiler_end(O_1, "test_target")
O_2 = relay.nn.relu(O_1)
ce_3 = compiler_end(O_2, "test_target")
cb_d = compiler_begin(ce_2, "default")
X = relay.tanh(cb_d)
ce_d = compiler_end(X, "default")
cb_3 = compiler_begin(ce_3, "test_target")
cb_4 = compiler_begin(ce_d, "test_target")
O_3 = relay.add(cb_3, cb_4)
O_4 = relay.add(cb_3, cb_4)
O_5 = relay.Tuple([O_3, O_4])
ce_4 = compiler_end(O_5, "test_target")
merged = relay.Function([data], ce_4)
region_set = relay.analysis.AnnotatedRegionSet(
merged, relay.op.get("annotation.compiler_begin"),
relay.op.get("annotation.compiler_end"))
assert len(region_set) == 3
check_region(
region_set,
"test_target",
[cb_1],
[cb_1, O_1, O_2, ce_2, ce_3],
[ce_2, ce_3],
)
check_region(
region_set,
"default",
[cb_d],
[cb_d, X, ce_d],
[ce_d],
)
check_region(
region_set,
"test_target",
[cb_3, cb_4],
[cb_3, cb_4, O_3, O_4, O_5, ce_4],
[ce_4],
)
def expected():
mod = tvm.IRModule()
# function 0
f0_i0 = relay.var(target + "_0_i0", shape=(10, 10))
f0_i1 = relay.var(target + "_0_i1")
f0_i2 = relay.var(target + "_0_i2")
f0_i3 = relay.var(target + "_0_i3")
f0_i4 = relay.var(target + "_0_i4")
f0_n0 = relay.nn.batch_norm(f0_i0, f0_i1, f0_i2, f0_i3, f0_i4)
f0_n1 = f0_n0[1]
f0_n2 = relay.nn.relu(f0_n0[0])
f0_o0 = relay.Tuple([f0_n2, f0_n1])
func0 = relay.Function([f0_i0, f0_i1, f0_i2, f0_i3, f0_i4], f0_o0)
func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Compiler", target)
func0 = func0.with_attr("global_symbol", target + "_0")
gv0 = relay.GlobalVar(target + "_0")
mod[gv0] = func0
mod = transform.InferType()(mod)
# body
data = relay.var("data", shape=(10, 10))
bn_gamma = relay.var("bn_gamma")
bn_beta = relay.var("bn_beta")
bn_mmean = relay.var("bn_mean")
bn_mvar = relay.var("bn_var")
function_out = gv0(data, bn_gamma, bn_beta, bn_mmean, bn_mvar)
get_out0 = relay.TupleGetItem(function_out, 0)
get_out1 = relay.TupleGetItem(function_out, 1)
out_2 = relay.tanh(get_out1)
out_3 = relay.log(get_out1)
out = relay.Tuple([get_out0, out_2, out_3])
func = relay.Function([data, bn_gamma, bn_beta, bn_mmean, bn_mvar], out)
mod["main"] = func
mod = transform.InferType()(mod)
return mod
def get_conv3d_transpose(
x_shape=(1, 32, 8, 8, 8),
k_shape=(32, 16, 3, 3, 3),
groups=1,
padding=(0, 0, 0),
strides=(1, 1, 1),
output_padding=(0, 0, 0),
activation=None,
dtype="float32",
data_layout="NCDHW",
kernel_layout="OIDHW",
):
x = relay.var("x", shape=(x_shape), dtype=dtype)
kernel = relay.const(np.random.randint(0, 1, k_shape).astype(dtype))
out = relay.nn.conv3d_transpose(
x,
kernel,
channels=k_shape[1],
kernel_size=k_shape[2:5],
groups=groups,
padding=padding,
strides=strides,
output_padding=output_padding,
data_layout=data_layout,
kernel_layout=kernel_layout,
)
dic = {"x": x_shape, "kernel": k_shape}
param_lst = ["kernel"]
if activation == "relu":
return relay.nn.relu(out), dic, param_lst
elif activation == "tanh":
return relay.tanh(out), dic, param_lst
elif activation == "sigmoid":
return relay.sigmoid(out), dic, param_lst
else:
return out, dic, param_lst
def expected():
mod = tvm.IRModule()
# function 1
f1_cb1 = relay.var("test_target_0_i0", shape=(10, 10))
f1_O_1 = relay.abs(f1_cb1)
f1_O_2 = relay.nn.relu(f1_O_1)
f1_out = relay.Tuple((f1_O_2, f1_O_1))
func1 = relay.Function([f1_cb1], f1_out)
func1 = set_func_attr(func1, "test_target", "test_target_0")
gv1 = relay.GlobalVar("test_target_0")
mod[gv1] = func1
mod = relay.transform.InferType()(mod)
# function 0
f2_cb3 = relay.var("test_target_1_i0", shape=(10, 10))
f2_cb4 = relay.var("test_target_1_i1", shape=(10, 10))
f2_O_3 = relay.add(f2_cb3, f2_cb4)
func0 = relay.Function([f2_cb3, f2_cb4], f2_O_3)
func0 = set_func_attr(func0, "test_target", "test_target_1")
gv0 = relay.GlobalVar("test_target_1")
mod[gv0] = func0
mod = relay.transform.InferType()(mod)
# body
data = relay.var("data", shape=(10, 10))
tuple_out = gv1(data)
ce_2 = relay.TupleGetItem(tuple_out, 1)
ce_3 = relay.TupleGetItem(tuple_out, 0)
X = relay.tanh(ce_2)
ce_4 = gv0(ce_3, X)
func = relay.Function([data], ce_4)
mod["main"] = func
mod = relay.transform.InferType()(mod)
return mod
def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""):
"""Long-Short Term Memory (LSTM) network cell.
Parameters
----------
num_hidden : int
Number of units in output symbol.
batch_size : int
Batch size (length of states).
Returns
-------
result : tvm.relay.Function
A Relay function that evaluates an LSTM cell.
The function takes in a tensor of input data, a tuple of two
states, and weights and biases for dense operations on the
inputs and on the state. It returns a tuple with two members,
an output tensor and a tuple of two new states.
"""
builder = relay.ScopeBuilder()
input_type = relay.TensorType((batch_size, num_hidden), dtype)
weight_type = relay.TensorType((4*num_hidden, num_hidden), dtype)
bias_type = relay.TensorType((4*num_hidden,), dtype)
dense_type = relay.TensorType((batch_size, 4*num_hidden), dtype)
slice_type = relay.TupleType([input_type, input_type,
input_type, input_type])
ret_type = relay.TupleType([input_type,
relay.TupleType([input_type, input_type])])
inputs = relay.Var("inputs", input_type)
states = relay.Var("states",
relay.TupleType([input_type, input_type]))
i2h_weight = relay.Var("i2h_weight", weight_type)
i2h_bias = relay.Var("i2h_bias", bias_type)
h2h_weight = relay.Var("h2h_weight", weight_type)
h2h_bias = relay.Var("h2h_bias", bias_type)
i2h = builder.let(("i2h", dense_type),
layers.dense_add_bias(
data=inputs,
units=num_hidden * 4,
weight=i2h_weight, bias=i2h_bias,
name="%si2h" % name))
h2h = builder.let(("h2h", dense_type),
layers.dense_add_bias(
data=relay.TupleGetItem(states, 0),
units=num_hidden * 4,
weight=h2h_weight, bias=h2h_bias,
name="%sh2h" % name))
gates = builder.let(("gates", dense_type), relay.add(i2h, h2h))
slice_gates = builder.let(("slice_gates", slice_type),
relay.split(gates,
indices_or_sections=4,
axis=1).astuple())
in_gate = builder.let(("in_gate", input_type),
relay.sigmoid(relay.TupleGetItem(slice_gates, 0)))
forget_gate = builder.let(("forget_gate", input_type),
relay.sigmoid(relay.TupleGetItem(slice_gates, 1)))
in_transform = builder.let(("in_transform", input_type),
relay.tanh(relay.TupleGetItem(slice_gates, 2)))
out_gate = builder.let(("out_gate", input_type),
relay.sigmoid(relay.TupleGetItem(slice_gates, 3)))
next_c = builder.let(("next_c", input_type),
relay.add(relay.multiply(forget_gate,
relay.TupleGetItem(states, 1)),
relay.multiply(in_gate, in_transform)))
next_h = builder.let(("next_h", input_type),
relay.multiply(out_gate, relay.tanh(next_c)))
ret = builder.let(("ret", ret_type),
relay.Tuple([next_h, relay.Tuple([next_h, next_c])]))
builder.ret(ret)
body = builder.get()
return relay.Function([inputs, states, i2h_weight,
i2h_bias, h2h_weight, h2h_bias],
body, ret_type)
def _execute(self):
self.node_dict = {}
# self.node_dict['1'] = relay.const(np.zeros((1, 128)), dtype='int32')
gelu_a = relay.var('gelu_a', shape=())
gelu_b = relay.var('gelu_b', shape=())
gelu_c = relay.var('gelu_c', shape=())
gelu_d = relay.var('gelu_d', shape=())
gelu_e = relay.var('gelu_e', shape=())
self.node_dict['1'] = relay.var('input.1', shape=(1,128), dtype='int32')
self.node_dict['2'] = relay.var('input.2', shape=(1,128), dtype='int32')
for gnode in self.graph:
name = gnode['name']
op_type = gnode['op_type']
attrs = gnode['attrs']
del attrs['A_shape']
del attrs['O_shape']
inputs = gnode['inputs']
if op_type == 'Const':
arr = np.zeros(attrs['shape'], dtype=np.int32)
y = relay.const(arr, dtype='int32')
elif op_type == 'expand_dims':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.expand_dims(x, attrs['axis'], attrs['num_newaxis'])
elif op_type == 'reshape':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.reshape(x, attrs['newshape'])
elif op_type == 'take':
data = get_input(self.node_dict, self.params, inputs[0])
indices = get_input(self.node_dict, self.params, inputs[1])
y = relay.take(data, indices, axis=attrs['axis'][0], mode=attrs['mode'])
elif op_type == 'one_hot':
x = get_input(self.node_dict, self.params, inputs[0])
cc1 = get_input(self.node_dict, self.params, inputs[1])
cc2 = get_input(self.node_dict, self.params, inputs[2])
y = relay.one_hot(x, cc1, cc2, **attrs)
elif op_type == 'strided_slice':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.strided_slice(x, **attrs)
elif op_type == 'mean':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.mean(x, axis=attrs['axis'],
exclude=attrs['exclude'],
keepdims=attrs['keepdims'])
elif op_type == 'nn.dense':
x = get_input(self.node_dict, self.params, inputs[0])
weight = get_input(self.node_dict, self.params, inputs[1])
y = relay.nn.dense(x, weight, units=attrs['units'][0])
elif op_type == 'add':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.add(x1, x2)
elif op_type == 'subtract':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.subtract(x1, x2)
elif op_type == 'multiply':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.multiply(x1, x2)
elif op_type == 'power':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.power(x1, x2)
elif op_type == 'transpose':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.transpose(x, **attrs)
elif op_type == 'tanh':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.tanh(x)
elif op_type == 'squeeze':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.squeeze(x, **attrs)
elif op_type == 'nn.batch_matmul':
x1 = get_input(self.node_dict, self.params, inputs[0])
x2 = get_input(self.node_dict, self.params, inputs[1])
y = relay.nn.batch_matmul(x1, x2)
elif op_type == 'nn.softmax':
x = get_input(self.node_dict, self.params, inputs[0])
y = relay.nn.softmax(x, **attrs)
elif op_type == 'gelu':
x = get_input(self.node_dict, self.params, inputs[0])
y = x * gelu_a * (gelu_b + relay.tanh(
( gelu_c * (x + gelu_d *
relay.power(x, gelu_e)))))
else:
import pdb; pdb.set_trace()
print( 'not supported op %s ' % op_type)
self.node_dict[name] = y
output_name = self.output_node_ids[0]
output = self.node_dict[output_name]
inputs = relay.analysis.free_vars(output)
# inputs = [self.node_dict['1'], self.node_dict['2']]
func = relay.Function(inputs, output)
mod = tvm.IRModule()
mod['main'] = func
with relay.build_config(opt_level=0):
graph, lib, params = relay.build(mod, 'llvm', params={})
self.m = graph_runtime.create(graph, lib, tvm.cpu())
def lstm_cell(num_hidden, batch_size=1, dtype="float32", name=""):
"""Long-Short Term Memory (LSTM) network cell.
Parameters
----------
num_hidden : int
Number of units in output symbol.
batch_size : int
Batch size (length of states).
Returns
-------
result : tvm.relay.Function
A Relay function that evaluates an LSTM cell.
The function takes in a tensor of input data, a tuple of two
states, and weights and biases for dense operations on the
inputs and on the state. It returns a tuple with two members,
an output tensor and a tuple of two new states.
"""
builder = relay.ScopeBuilder()
input_type = relay.TensorType((batch_size, num_hidden), dtype)
weight_type = relay.TensorType((4 * num_hidden, num_hidden), dtype)
bias_type = relay.TensorType((4 * num_hidden,), dtype)
dense_type = relay.TensorType((batch_size, 4 * num_hidden), dtype)
slice_type = relay.TupleType([input_type, input_type, input_type, input_type])
ret_type = relay.TupleType([input_type, relay.TupleType([input_type, input_type])])
inputs = relay.Var("inputs", input_type)
states = relay.Var("states", relay.TupleType([input_type, input_type]))
i2h_weight = relay.Var("i2h_weight", weight_type)
i2h_bias = relay.Var("i2h_bias", bias_type)
h2h_weight = relay.Var("h2h_weight", weight_type)
h2h_bias = relay.Var("h2h_bias", bias_type)
i2h = builder.let(
("i2h", dense_type),
layers.dense_add_bias(
data=inputs, units=num_hidden * 4, weight=i2h_weight, bias=i2h_bias, name="%si2h" % name
),
)
h2h = builder.let(
("h2h", dense_type),
layers.dense_add_bias(
data=relay.TupleGetItem(states, 0),
units=num_hidden * 4,
weight=h2h_weight,
bias=h2h_bias,
name="%sh2h" % name,
),
)
gates = builder.let(("gates", dense_type), relay.add(i2h, h2h))
slice_gates = builder.let(
("slice_gates", slice_type), relay.split(gates, indices_or_sections=4, axis=1).astuple()
)
in_gate = builder.let(
("in_gate", input_type), relay.sigmoid(relay.TupleGetItem(slice_gates, 0))
)
forget_gate = builder.let(
("forget_gate", input_type), relay.sigmoid(relay.TupleGetItem(slice_gates, 1))
)
in_transform = builder.let(
("in_transform", input_type), relay.tanh(relay.TupleGetItem(slice_gates, 2))
)
out_gate = builder.let(
("out_gate", input_type), relay.sigmoid(relay.TupleGetItem(slice_gates, 3))
)
next_c = builder.let(
("next_c", input_type),
relay.add(
relay.multiply(forget_gate, relay.TupleGetItem(states, 1)),
relay.multiply(in_gate, in_transform),
),
)
next_h = builder.let(("next_h", input_type), relay.multiply(out_gate, relay.tanh(next_c)))
ret = builder.let(("ret", ret_type), relay.Tuple([next_h, relay.Tuple([next_h, next_c])]))
builder.ret(ret)
body = builder.get()
return relay.Function(
[inputs, states, i2h_weight, i2h_bias, h2h_weight, h2h_bias], body, ret_type
)
def tvm_elem_tanh(node, ctx):
inp = ctx[node.args[0].name]
th = relay.tanh(inp)
return node.args[1].name, th
def get_inner_func_1():
x = relay.var("x", shape=(1, 4, 5, 6), dtype="int8")
x = relay.tanh(x)
x = _create_primitive_function(x)
return x
