def before(x):
inj = relay.squeeze(x)
y1 = relay.add(inj, relay.const(1, "float32"))
tmp = relay.squeeze(inj)
tmp = relay.add(tmp, relay.const(1, "float32"))
y2 = relay.add(tmp, relay.const(1, "float32"))
y3 = relay.add(inj, relay.const(1, "float32"))
concat = relay.concatenate((y1, y2, y3), axis=1)
out_inj = relay.squeeze(concat)
out = relay.add(out_inj, relay.const(1, "float32"))
return relay.Function(relay.ir_pass.free_vars(out), out)
def before():
x = relay.var("x", shape=(10, 20))
y = relay.add(x, relay.const(1, "float32"))
z = relay.squeeze(y)
u = relay.transpose(y, axes=[0, 1])
w = relay.left_shift(z, u)
return relay.Function([x], w)
def test_squeeze_infer_type():
n, t, d = 1, 4, 1
x = relay.var("x", relay.TensorType((n, t, d), "float32"))
y = relay.squeeze(x, axis=(2,))
assert "axis=" in y.astext()
yy = relay.ir_pass.infer_type(y)
assert yy.checked_type == relay.TensorType(
(1, 4), "float32")
n, t, d = 1, 4, 1
x = relay.var("x", relay.TensorType((n, t, d), "float32"))
y = relay.squeeze(x)
assert "axis=" not in y.astext()
yy = relay.ir_pass.infer_type(y)
assert yy.checked_type == relay.TensorType(
(4,), "float32")
def expected():
x = relay.var("p", shape=(10, 20))
y = relay.add(x, relay.const(1, "float32"))
z = relay.exp(y)
w = relay.squeeze(z)
f1 = relay.Function([x], w)
x = relay.var("x", shape=(10, 20))
y = relay.Call(f1, [x])
return relay.Function([x], y)
def expected():
x = relay.var("p", shape=(10, 20))
y = relay.add(x, relay.const(1, "float32"))
z = relay.squeeze(y)
u = relay.transpose(y, axes=[0, 1])
w = relay.left_shift(z, u)
f1 = relay.Function([x], w)
x = relay.var("x", shape=(10, 20))
y = relay.Call(f1, [x])
return relay.Function([x], y)
def expected(x, conv_weight, out_scale, channels):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight]
squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
conv_weight = relay.multiply(
conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
padding=(1, 1))
return relay.Function(args, y)
def verify_squeeze(shape, dtype, axis):
x = relay.var("x", relay.TensorType(shape, dtype))
squeeze = relay.squeeze(x, axis=axis)
np_axis = tuple(axis) if axis is not None else None
data = np.random.random_sample(shape).astype(dtype)
intrp = create_executor()
op_res = intrp.evaluate(squeeze, { x : relay.const(data) })
ref_res = np.squeeze(data, axis=np_axis)
np.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=0.01)
def expected(x, conv_weight, in_bias, in_scale, channels):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, in_bias]
in_bias = relay.expand_dims(in_bias, axis=1, num_newaxis=2)
squeezed_scale = relay.squeeze(in_scale, axis=[1,2])
x = relay.nn.relu(x)
in_bias = relay.divide(in_bias, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
x = relay.add(x, in_bias)
conv_weight = relay.multiply(
conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
padding=(1, 1))
return relay.Function(args, y)
def expected(x, conv_weight, out_scale, channels, blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight]
if blocking:
squeezed_scale = relay.squeeze(out_scale, axis=[0, 2, 3])
conv_weight = relay.multiply(
conv_weight,
relay.reshape(squeezed_scale, (channels // blocking[1], 1, 1, 1, 1, blocking[1])),
)
else:
squeezed_scale = relay.squeeze(out_scale, axis=[1, 2])
conv_weight = relay.multiply(
conv_weight, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3)
)
y = relay.nn.conv2d(
x,
conv_weight,
channels=channels,
kernel_size=(3, 3),
padding=(1, 1),
data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW",
)
return relay.Function(args, y)
def expected(x, conv_weight, in_bias, in_scale, channels):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, in_bias, in_scale]
in_scale = relay.expand_dims(in_scale, axis=1, num_newaxis=2)
in_bias = relay.expand_dims(in_bias, axis=1, num_newaxis=2)
squeezed_scale = relay.squeeze(in_scale, axis=[1,2])
x = relay.nn.relu(x)
in_bias = relay.divide(in_bias, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
x = relay.add(x, in_bias)
conv_weight = relay.multiply(
conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
padding=(1, 1))
return relay.Function(args, y)
def expected(x, conv_weight, out_bias, out_scale, in_channels, channels,
blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, out_bias]
if not blocking:
out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
squeezed_scale = relay.squeeze(out_scale, axis=[1, 2])
def fold_conv_weight():
if blocking:
return relay.multiply(
conv_weight,
relay.reshape(
squeezed_scale,
(channels // blocking[1], 1, 1, 1, 1, blocking[1])),
)
else:
return relay.multiply(
conv_weight,
relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y1 = relay.nn.conv2d(
x,
fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
padding=(1, 1),
data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
kernel_layout="OIHW1i{}o".format(blocking[1])
if blocking else "OIHW",
)
y1 = relay.nn.relu(y1)
y2 = relay.nn.conv2d(
x,
fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
padding=(1, 1),
data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
kernel_layout="OIHW1i{}o".format(blocking[1])
if blocking else "OIHW",
)
y2 = relay.nn.relu(y2)
y = relay.add(y1, y2)
return relay.Function(args, y)
def build_relay_module(batch_size, input_size, hidden_size, time_steps, dense_dim):
mod = tvm.IRModule()
mod["lstm_layer"] = lstm_definition(batch_size, input_size, hidden_size, time_steps)
mod["linear_layer"] = linear_layer_definition(batch_size, hidden_size, dense_dim)
lstm_var = mod.get_global_var("lstm_layer")
linear_var = mod.get_global_var("linear_layer")
# now we build up our main function
input_var = relay.var("input", shape=(batch_size, time_steps, input_size))
init_hidden_var = relay.var("init_hidden", shape=(batch_size, hidden_size))
init_cell_var = relay.var("init_cell", shape=(batch_size, hidden_size))
i2h_weight_var = relay.var("i2h_weight", shape=(4*hidden_size, input_size))
h2h_weight_var = relay.var("h2h_weight", shape=(4*hidden_size, hidden_size))
lstm_bias_var = relay.var("lstm_bias", shape=(4*hidden_size,))
linear_weight_var = relay.var("linear_weight", shape=(dense_dim, hidden_size))
linear_bias_var = relay.var("linear_bias", shape=(dense_dim,))
builder = relay.ScopeBuilder()
state_var = builder.let("state", relay.Tuple([init_hidden_var, init_cell_var]))
lstm_res = builder.let("lstm_res",
lstm_var(input_var, state_var,
i2h_weight_var, h2h_weight_var,
lstm_bias_var,
# the keras model only gave one bias,
# so set the other to zero
# (hopefully this is correct)
relay.zeros_like(lstm_bias_var)))
final_hidden = builder.let("final_hidden",
relay.TupleGetItem(lstm_res, 1))
# to match PT's semantics, we're undoing the reshape in LSTM :)
reshape_hidden = builder.let("reshape_hidden",
relay.squeeze(final_hidden, axis=[0]))
linear_result = builder.let("linear_result",
linear_var(reshape_hidden,
linear_weight_var, linear_bias_var))
# finally do a softmax
builder.ret(relay.nn.softmax(linear_result))
main_func = relay.Function([input_var, init_hidden_var, init_cell_var,
i2h_weight_var, h2h_weight_var, lstm_bias_var,
linear_weight_var, linear_bias_var],
builder.get())
mod["main"] = main_func
return mod
def verify_squeeze(shape, axis, oshape):
x = relay.var("x", relay.TensorType(shape, "float32"))
y = relay.var("y", relay.TensorType(axis, "float32"))
z = relay.squeeze(x, relay.shape_of(y))
func = run_infer_type(relay.Function([x, y], z))
func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()),
transform.InferType())
zz = func2.body
assert isinstance(zz, relay.Call)
assert zz.op == relay.op.get("squeeze")
assert "axis=" in zz.astext()
assert zz.checked_type == relay.ty.TensorType(oshape, "float32")
x_data = np.random.uniform(low=-1, high=1,
size=shape).astype("float32")
y_data = np.random.uniform(low=-1, high=1, size=axis).astype("float32")
ref_res = np.squeeze(x_data, axis)
verify_func(func2, [x_data, y_data], ref_res)
def expected(x, conv_weight, out_bias, out_scale, channels):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, out_bias]
out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
def fold_conv_weight():
return relay.multiply(
conv_weight ,
relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y1 = relay.nn.conv2d(x, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
padding=(1, 1))
y1 = relay.nn.relu(y1)
y2 = relay.nn.conv2d(x, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
padding=(1, 1))
y2 = relay.nn.relu(y2)
y = relay.add(y1, y2)
return relay.Function(args, y)
def expected(x, conv_weight, out_bias, out_scale, channels):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, out_bias]
out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
def fold_conv_weight():
return relay.multiply(
conv_weight ,
relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y1 = relay.nn.conv2d(x, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
padding=(1, 1))
y1 = relay.nn.relu(y1)
y2 = relay.nn.conv2d(x, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
padding=(1, 1))
y2 = relay.nn.relu(y2)
y = relay.add(y1, y2)
return relay.Function(args, y)
def expected(x, conv_weight, out_bias, out_scale, channels):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, out_bias]
squeezed_scale = relay.squeeze(out_scale, axis=[1, 2])
conv_weight = relay.multiply(
conv_weight,
relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y = relay.nn.conv2d(
x,
conv_weight,
channels=channels,
kernel_size=(3, 3),
padding=(1, 1),
data_layout="NCHW",
kernel_layout="OIHW",
)
out_bias = relay.multiply(out_bias, squeezed_scale)
y = relay.nn.bias_add(y, out_bias)
y = relay.nn.relu(y)
return relay.Function(args, y)
def before():
x = relay.var("x", shape=(10, 20))
y = relay.add(x, relay.const(1, "float32"))
z = relay.exp(y)
w = relay.squeeze(z)
return relay.Function([x], w)
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 get_graph(x_shape, axis):
x = relay.var("x", shape=(x_shape), dtype="float32")
out = relay.squeeze(x, axis=axis)
f = relay.Function([x], out)
return f, {"x": x_shape}, []
def fold_conv_weight():
squeezed_scale = relay.squeeze(out_scale, axis=[1, 2])
return relay.multiply(
conv_weight,
relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
def test_squeeze_grad():
data = relay.var("data", shape=(2, 1, 1, 3, 4, 1), dtype="float64")
fwd_func = relay.Function([data], relay.squeeze(data))
fwd_func_subset = relay.Function([data], relay.squeeze(data, axis=[1, -1]))
check_grad(fwd_func)
check_grad(fwd_func_subset)
def lstm_definition(batch_size,
input_size,
hidden_size,
time_steps,
time_axis=1):
state_tensor_type = relay.TensorType((batch_size, hidden_size))
state_tuple_type = relay.TupleType([state_tensor_type, state_tensor_type])
input_var = relay.var("input", shape=(batch_size, time_steps, input_size))
state_var = relay.var("state", type_annotation=state_tuple_type)
i2h_weight_var = relay.var("i2h_weight",
shape=(4 * hidden_size, input_size))
h2h_weight_var = relay.var("h2h_weight",
shape=(4 * hidden_size, hidden_size))
i2h_bias_var = relay.var("i2h_bias", shape=(4 * hidden_size, ))
h2h_bias_var = relay.var("h2h_bias", shape=(4 * hidden_size, ))
# in this case, we are ignoring the state outputs
builder = relay.ScopeBuilder()
cell_var = builder.let(
"lstm_cell", relay_lstm_cell(batch_size, input_size, hidden_size))
splits = builder.let(
"splits",
relay.split(input_var, time_steps, time_axis).astuple())
last_state = state_var
seq_outs = []
for i in range(time_steps):
squeezed = builder.let(
f"squeezed_{i}",
relay.squeeze(relay.TupleGetItem(splits, i), axis=[time_axis]))
cell_out = builder.let(
f"cell_out_{i}",
cell_var(squeezed, last_state, i2h_weight_var, h2h_weight_var,
i2h_bias_var, i2h_bias_var))
new_seq_out = builder.let(f"seq_out_{i}",
relay.TupleGetItem(cell_out, 0))
seq_outs.append(new_seq_out)
new_hidden = builder.let(f"state_update_{i}",
relay.TupleGetItem(cell_out, 1))
last_state = new_hidden
stacked = builder.let("stacked", relay.stack(seq_outs, axis=time_axis))
# finally reshape to match pytorch's semantics (one layer)
reshape_hidden = builder.let(
"final_hidden",
relay.reshape(relay.TupleGetItem(last_state, 0),
(1, batch_size, hidden_size)))
reshape_cell = builder.let(
"final_cell",
relay.reshape(relay.TupleGetItem(last_state, 1),
(1, batch_size, hidden_size)))
builder.ret(relay.Tuple([stacked, reshape_hidden, reshape_cell]))
ret_type = relay.TupleType([
relay.TensorType((batch_size, time_steps, hidden_size)),
relay.TensorType((1, batch_size, hidden_size)),
relay.TensorType((1, batch_size, hidden_size))
])
return relay.Function([
input_var, state_var, i2h_weight_var, h2h_weight_var, i2h_bias_var,
h2h_bias_var
],
builder.get(),
ret_type=ret_type)
def test_squeeze_bad_axes_infer_type():
n, t, d = 1, 4, 1
x = relay.var("x", relay.TensorType((n, t, d), "float32"))
y = relay.squeeze(x, axis=(1,))
yy = relay.ir_pass.infer_type(y)
def fold_conv_weight():
squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
return relay.multiply(
conv_weight ,
relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
def test_squeeze_bad_axes_infer_type():
(n, t, d) = (1, 4, 1)
x = relay.var('x', relay.TensorType((n, t, d), 'float32'))
y = relay.squeeze(x, axis=(1,))
yy = run_infer_type(y)
def before():
x = relay.var("x", shape=(10, 20))
y = relay.add(x, relay.const(1, "float32"))
z = relay.exp(y)
w = relay.squeeze(z)
return relay.Function([x], w)
def verify_squeeze(shape, dtype, axis):
x = relay.var("x", relay.TensorType(shape, dtype))
z = relay.squeeze(x, axis=axis)
func = relay.Function([x], z)
x_data = np.random.random_sample(shape).astype(dtype)
verify_results(func, [x_data], "test_squeeze", rtol=1e-5, atol=1e-5)
def test_squeeze_bad_axes_infer_type():
n, t, d = 1, 4, 1
x = relay.var("x", relay.TensorType((n, t, d), "float32"))
y = relay.squeeze(x, axis=(1,))
yy = relay.ir_pass.infer_type(y)
def test_squeeze(x_shape, axis):
x = relay.var('x', shape=(x_shape), dtype='float32')
out = relay.squeeze(x, axis=axis)
f = relay.Function([x], out)
return f, {'x': x_shape}
