def custom_lstm_test():
input_name = 'X'
seq_len = 5
batch = 2
input_size = 3
hidden_size = 4
num_layers = 4
input_shapes = {}
input_types = {
input_name:
TensorType((seq_len, batch, input_size)),
"states":
TupleType([
TensorType((batch, hidden_size)),
TensorType((batch, hidden_size))
])
}
from custom_lstms import rnn_layer, stacked_rnn, stacked_lnlstm
models = [
rnn_layer(input_size, hidden_size).eval(),
stacked_rnn(input_size, hidden_size, num_layers).eval(),
stacked_lnlstm(input_size, hidden_size, num_layers).eval()
]
for raw_model in models:
script_module = torch.jit.script(raw_model)
mod, params = from_pytorch(script_module, input_shapes, input_types)
# comp = relay.backend.vm.VMCompiler()
# opt_mod, _ = comp.optimize(mod, "llvm", params)
# print(opt_mod["main"])
# continue
for i in range(5):
inp = torch.randn(seq_len, batch, input_size)
states = [(torch.randn(batch, hidden_size),
torch.randn(batch, hidden_size))
for _ in range(num_layers)]
with torch.no_grad():
pt_result = raw_model(inp.clone(), states[0])
params[input_name] = inp.numpy()
params["states"] = (st.numpy() for st in states[0])
run_and_compare(mod, params, pt_result)
def build_impl(self, input_size, memory_size, dtype="float32"):
t = TensorType(shape=(1, memory_size), dtype=dtype)
self.ret_type = TupleType([t, t])
tree_type = self.p.tree(TensorType(shape=(1, input_size), dtype=dtype))
t = self.input(Var("tlstm_input", tree_type))
i = Var("i", TensorType(shape=(1, input_size), dtype=dtype))
c = Var("c", self.p.l(tree_type))
cell = LSTMCell(input_size=input_size,
memory_size=memory_size,
dtype=dtype)
rose_case = Clause(
PatternConstructor(self.p.rose,
[PatternVar(i), PatternVar(c)]),
cell(i, self.p.map(lam(["x"], self), c)))
return Match(t, [rose_case])
def build_impl(self, input_size, memory_size, dtype="float32"):
t = TensorType(shape=(1, memory_size), dtype=dtype)
i = self.input(
var("lstmcell_input", shape=(1, input_size), dtype=dtype))
c = self.input(Var("lstmcell_children", self.p.l(TupleType([t, t]))))
sum = lam(["x", "y"], lambda x, y: x + y)
child_h_sum = self.p.foldl(
sum, op.zeros(shape=(1, memory_size), dtype=dtype),
self.p.map(lam(["z"], lambda z: TupleGetItem(z, 1)), c))
ioux = Linear(input_size=input_size, output_size=memory_size * 3)(i)
iouh = Linear(input_size=memory_size,
output_size=memory_size * 3)(child_h_sum)
iou = ioux + iouh
fx = Linear(input_size=input_size, output_size=memory_size)(i)
fh = Linear(input_size=memory_size, output_size=memory_size)
i, o, u = op.split(iou, 3, axis=1)
i, o, u = op.sigmoid(i), op.sigmoid(o), op.tanh(u)
def foreach_children(children):
f = op.sigmoid(fh(TupleGetItem(children, 1)) + fx)
return f * TupleGetItem(children, 0)
c = self.p.foldl(sum, i * u, self.p.map(lam(["z"], foreach_children),
c))
return Tuple([c, o * op.tanh(c)])
def check_grad(func, mod=None):
"""
Test that directional gradient calculated by reverse mode
is close to the one calculated by finite difference.
"""
global CHECK_GRAD_COUNTER
if mod is None:
mod = relay.Module()
def make(name):
return GlobalVar(name + str(CHECK_GRAD_COUNTER))
func_name = make("func_")
back_func_name = make("back_func_")
finite_difference_func_name = make("finite_difference_")
reverse_mode_func_name = make("reverse_mode_")
check_func_name = make("check_func_")
CHECK_GRAD_COUNTER = CHECK_GRAD_COUNTER + 1
epsilon = relay.const(0.01)
mod[func_name] = func
mod[back_func_name] = gradient(mod[func_name], mod=mod)
params = mod[func_name].params
directions = [rand_from_type(x.checked_type) for x in params]
ft = TensorType(())
sb = ScopeBuilder()
def get_reverse_mode_result(e, d, t):
assert isinstance(t, TensorType)
return op.cast(e * d, 'float32')
bf = sb.let("bf", TupleGetItem(back_func_name(*params), 1))
reverse_mode_results = [
get_reverse_mode_result(TupleGetItem(bf, i), directions[i],
x.checked_type) for i, x in enumerate(params)
]
reverse_mode_result = relay.const(0.0)
for x in reverse_mode_results:
reverse_mode_result = reverse_mode_result + op.reduce.sum(x)
sb.ret(reverse_mode_result)
reverse_mode_result = sb.get()
mod[reverse_mode_func_name] = Function(params, reverse_mode_result, ft,
mod[func_name].type_params,
mod[func_name].attrs)
finite_difference_result = op.reduce.sum(
(func_name(*[x + epsilon * y for x, y in zip(params, directions)]) -
func_name(*params)) / epsilon)
mod[finite_difference_func_name] = Function(params,
finite_difference_result, ft,
mod[func_name].type_params,
mod[func_name].attrs)
check_func_result = op.abs(
reverse_mode_func_name(*params) - finite_difference_func_name(*params))
mod[check_func_name] = Function(params, check_func_result, ft,
mod[func_name].type_params,
mod[func_name].attrs)
ex = create_executor(mod=mod)
res = ex.evaluate(
check_func_name(*[rand_from_type(x.checked_type) for x in params]))
assert res.data.asnumpy() < 0.001
def test_empty_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d)
g = dcpe(gradient(f))
expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])]))
assert alpha_equal(g, expected)
def test_empty_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d)
g = dcpe(f, grad=True)
expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])]))
expected = run_opt_pass(expected, transform.InferType())
assert tvm.ir.structural_equal(g, expected)
def test_empty_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d)
g = dcpe(f, grad=True)
expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])]))
expected = transform.OptimizeOnExpr(expected, transform.InferType())
assert alpha_equal(g, expected)
def test_empty_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d)
# TODO(mbs): Revisit once DCE eliminates dead writes.
g = dcpe(f, grad=True, ignore_impurity=True)
expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])]))
expected = run_opt_pass(expected, transform.InferType())
assert tvm.ir.structural_equal(g, expected)
def build_impl(self, input_size, memory_size, dtype="float32"):
l = self.input(
Var("l", self.p.l(TensorType(shape=(1, input_size), dtype=dtype))))
cell = LSTMCell(input_size=input_size,
memory_size=memory_size,
dtype=dtype)
return self.p.foldl(
lam(["c", "x"],
lambda c, x: cell(x, self.p.cons(c, self.p.nil()))),
Tuple([
op.zeros(shape=(1, memory_size), dtype=dtype),
op.zeros(shape=(1, memory_size), dtype=dtype)
]), l)
def build_impl(self, input_size, memory_size, dtype="float32"):
l = self.input(
Var("l", self.p.l(TensorType(shape=(1, input_size), dtype=dtype))))
def LSTM(l):
return LSTMTransformer(input_size=input_size,
memory_size=memory_size,
dtype=dtype)(l)
fwd = LSTM(l)
rev = LSTM(self.p.rev(l))
lhs = op.concatenate(
[TupleGetItem(fwd, 0), TupleGetItem(rev, 0)], axis=1)
t = TensorType(shape=(1, memory_size), dtype=dtype)
x = Var("x", TupleType([t, t])) # cannot infer here
rhs = self.p.map(
Function([x],
op.concatenate([TupleGetItem(x, 0),
TupleGetItem(x, 1)],
axis=1)),
self.p.zip(TupleGetItem(fwd, 1), TupleGetItem(rev, 1)))
return Tuple([lhs, rhs])
def test_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d * d)
g = dcpe(f, grad=True)
m = d * d
x = relay.Var("x")
o = op.ones_like(x)
x1 = relay.Var("x1")
grad = op.zeros_like(d) + op.collapse_sum_like(x1 * d, d) + op.collapse_sum_like(x1 * d, d)
body = Tuple([x, Tuple([grad])])
body = relay.Let(x1, o, body)
expected = Function([d], relay.Let(x, m, body))
expected = run_opt_pass(expected, transform.InferType())
assert_alpha_equal(g, expected)
def test_if_ref():
shape = ()
dtype = "bool"
t = TensorType(shape, dtype)
d = Var("d", t)
r = Var("r")
update = Function([], RefWrite(r, RefRead(r) + RefRead(r)))
u = Var("u")
body = If(d, u(), u())
eff = Var("eff")
body = Let(eff, body, RefRead(r))
f = Function([d], Let(r, RefCreate(const(1)), Let(u, update, body)))
pe_f = tipe(f)
f_res = create_executor().evaluate(f)(const(True))
pe_f_res = create_executor().evaluate(pe_f)(const(True))
np.testing.assert_allclose(f_res.numpy(), 2 * np.ones_like(f_res.numpy()))
np.testing.assert_allclose(pe_f_res.numpy(), 2 * np.ones_like(pe_f_res.numpy()))
def test_ad():
shape = (10, 10)
dtype = "float32"
t = TensorType(shape, dtype)
d = Var("d", t)
f = Function([d], d * d)
# TODO(mbs): Revisit once DCE eliminates dead writes.
g = dcpe(f, grad=True, ignore_impurity=True)
m = d * d
x = relay.Var("x")
o = op.ones_like(x)
x1 = relay.Var("x1")
grad = op.zeros_like(d) + op.collapse_sum_like(x1 * d, d) + op.collapse_sum_like(x1 * d, d)
body = Tuple([x, Tuple([grad])])
body = relay.Let(x1, o, body)
expected = Function([d], relay.Let(x, m, body))
expected = run_opt_pass(expected, transform.InferType())
tvm.ir.assert_structural_equal(g, expected)
def build_impl(self, input_size, memory_size, dtype="float32"):
l = self.input(
Var("l", self.p.l(TensorType(shape=(1, input_size), dtype=dtype))))
def f(c, x):
cell = LSTMCell(input_size=input_size,
memory_size=memory_size,
dtype=dtype)
o = cell(x, self.p.cons(c, self.p.nil()))
return Tuple([o, TupleGetItem(o, 1)])
res = self.p.map_accuml(
lam(["c", "x"], f),
Tuple([
op.zeros(shape=(1, memory_size), dtype=dtype),
op.zeros(shape=(1, memory_size), dtype=dtype)
]), l)
return Tuple(
[TupleGetItem(TupleGetItem(res, 0), 1),
TupleGetItem(res, 1)])
def test_function_invalidate():
shape = ()
dtype = "bool"
t = TensorType(shape, dtype)
d = Var("d", t)
r = Var("r")
fetch = Function([], RefRead(r))
fet = Var("fetch")
fet_obscured = Var("fetch_obscured")
u = Var("u")
body = If(d, fet_obscured(), fet_obscured())
body = Let(u, RefWrite(r, const(1)), body)
body = Let(fet_obscured, If(d, fet, fet), body)
body = Let(fet, fetch, body)
body = Let(r, RefCreate(const(0)), body)
f = Function([d], body)
pe_f = tipe(f)
f_res = create_executor().evaluate(f)(const(True))
pe_f_res = create_executor().evaluate(pe_f)(const(True))
np.testing.assert_allclose(f_res.numpy(), np.ones_like(f_res.numpy()))
np.testing.assert_allclose(pe_f_res.numpy(), np.ones_like(pe_f_res.numpy()))
