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()))