def run_on_models(models, inputs, target="llvm"): for raw_model in models: with torch.no_grad(): pt_result = raw_model(*inputs).numpy() script_module = torch.jit.trace(raw_model, *inputs).eval() input_names = get_graph_input_names(script_module) input_shapes = dict(zip(input_names, [inp.shape for inp in inputs])) mod, params = from_pytorch(script_module, input_shapes) with relay.build_config(opt_level=3): json, lib, params = relay.build(mod, target=target, params=params) ctx = tvm.context(target, 0) runtime = tvm.contrib.graph_runtime.create(json, lib, ctx) runtime.set_input(**params) for name, inp in zip(input_names, inputs): runtime.set_input(name, inp.numpy()) runtime.run() tvm_result = runtime.get_output(0).asnumpy() np.allclose(tvm_result, pt_result, rtol=1e-5, atol=1e-5) print(np.max(np.abs(tvm_result - pt_result)), np.mean(np.abs(tvm_result - pt_result))) tvm.testing.assert_allclose(tvm_result, pt_result, rtol=1e-3, atol=1e-3)
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 simple_rnn_test(): class DecisionGate(torch.nn.Module): def forward(self, x): if x.sum() > 0: return x else: return -x class Cell(torch.nn.Module): def __init__(self, dg): super(Cell, self).__init__() self.dg = dg self.linear = torch.nn.Linear(4, 4) def forward(self, x, h): new_h = torch.tanh(self.dg(self.linear(x)) + h) return new_h, new_h class RNNLoop(torch.nn.Module): def __init__(self): super().__init__() x = torch.rand(10, 4, dtype=torch.float) h = torch.rand(10, 4, dtype=torch.float) self.cell = torch.jit.trace(Cell(DecisionGate()), (x, h)) def forward(self, xs): h = torch.zeros(10, 4, dtype=torch.float) y = torch.zeros(10, 4, dtype=torch.float) for i in range(xs.size(0)): y, h = self.cell(xs[i], h) return y raw_model = RNNLoop() script_module = torch.jit.script(raw_model) input_name = get_graph_input_names(script_module)[0] input_shapes = {input_name: (10, 10, 4)} mod, params = from_pytorch(script_module, input_shapes, {}) for i in range(5): inp = torch.randn(input_shapes[input_name], dtype=torch.float) with torch.no_grad(): pt_result = raw_model(inp.clone()) params[input_name] = inp.numpy() run_and_compare(mod, params, pt_result)
def test_custom_lstm(): input_name = "input" states_name = "states" seq_len = 5 batch = 2 input_size = 3 hidden_size = 4 num_layers = 3 state_tensor_shape = (batch, hidden_size) torch.manual_seed(1) inp = torch.randn(seq_len, batch, input_size) input_shapes = [ (input_name, (seq_len, batch, input_size)), (states_name, (state_tensor_shape, state_tensor_shape)), ] input_shapes_stacked = [ (input_name, (seq_len, batch, input_size)), ( states_name, [(state_tensor_shape, state_tensor_shape), (state_tensor_shape, state_tensor_shape)], ), ] input_shapes_stacked_bidir = [ (input_name, (seq_len, batch, input_size)), ( states_name, [ [(state_tensor_shape, state_tensor_shape) for _ in range(2)] for _ in range(num_layers) ], ), ] states = [ (torch.randn(state_tensor_shape), torch.randn(state_tensor_shape)) for _ in range(num_layers) ] bidir_states = [ (torch.randn(state_tensor_shape), torch.randn(state_tensor_shape)) for _ in range(2) ] stacked_bidir_states = [ [(torch.randn(state_tensor_shape), torch.randn(state_tensor_shape)) for _ in range(2)] for _ in range(num_layers) ] models = [ ("lstm", lstm(input_size, hidden_size).eval(), states[0], input_shapes), ( "stacked", stacked_lstm(input_size, hidden_size, num_layers).eval(), states, input_shapes_stacked, ), ("bidir", bidir_lstm(input_size, hidden_size).eval(), bidir_states, input_shapes_stacked), # TODO(masahi): stacked bidir seems to have a rare accuracy issue # ( # "stacked_bidir", # stacked_bidir_lstm(input_size, hidden_size, num_layers).eval(), # stacked_bidir_states, # input_shapes_stacked_bidir, # ), ] for (name, raw_model, states, input_shapes) in models: script_module = torch.jit.script(raw_model) mod, params = from_pytorch(script_module, input_shapes) with torch.no_grad(): pt_result = raw_model(inp.clone(), states) params[input_name] = inp.numpy() if isinstance(states, tuple): states_np = tuple(st.numpy() for st in states) elif isinstance(states, list) and isinstance(states[0], torch.Tensor): states_np = [st.numpy() for st in states] elif isinstance(states, list) and isinstance(states[0], tuple): states_np = [tuple(st.numpy() for st in states[i]) for i in range(len(states))] elif isinstance(states, list) and isinstance(states[0], list): states_np = [ [tuple(st.numpy() for st in states) for states in states[layer]] for layer in range(num_layers) ] else: assert False if isinstance(states_np, list): params[states_name] = convert_list_to_vmobj(states_np) else: params[states_name] = states_np for tgt, ctx in tvm.testing.enabled_targets(): print("Running %s on target %s" % (name, tgt)) run_and_compare(mod, params, pt_result, target=tgt, ctx=ctx)
def custom_lstm_test(): input_name = "input" states_name = "states" seq_len = 5 batch = 2 input_size = 3 hidden_size = 4 num_layers = 3 state_tensor_shape = (batch, hidden_size) inp = torch.randn(seq_len, batch, input_size) input_shapes = [(input_name, (seq_len, batch, input_size)), (states_name, (state_tensor_shape, state_tensor_shape))] input_shapes_stacked = [(input_name, (seq_len, batch, input_size)), (states_name, [(state_tensor_shape, state_tensor_shape), (state_tensor_shape, state_tensor_shape)])] input_shapes_stacked_bidir = [ (input_name, (seq_len, batch, input_size)), (states_name, [[(state_tensor_shape, state_tensor_shape) for _ in range(2)] for _ in range(num_layers)]) ] states = [(torch.randn(state_tensor_shape), torch.randn(state_tensor_shape)) for _ in range(num_layers)] bidir_states = [(torch.randn(state_tensor_shape), torch.randn(state_tensor_shape)) for _ in range(2)] stacked_bidir_states = [[(torch.randn(state_tensor_shape), torch.randn(state_tensor_shape)) for _ in range(2)] for _ in range(num_layers)] models = [(lstm(input_size, hidden_size).eval(), states[0], input_shapes), (stacked_lstm(input_size, hidden_size, num_layers).eval(), states, input_shapes_stacked), (bidir_lstm(input_size, hidden_size).eval(), bidir_states, input_shapes_stacked), (stacked_bidir_lstm(input_size, hidden_size, num_layers).eval(), stacked_bidir_states, input_shapes_stacked_bidir)] for (raw_model, states, input_shapes) in models: script_module = torch.jit.script(raw_model) mod, params = from_pytorch(script_module, input_shapes) with torch.no_grad(): pt_result = raw_model(inp.clone(), states) params[input_name] = inp.numpy() if isinstance(states, tuple): states_np = tuple(st.numpy() for st in states) elif isinstance(states, list) and isinstance(states[0], torch.Tensor): states_np = [st.numpy() for st in states] elif isinstance(states, list) and isinstance(states[0], tuple): states_np = [ tuple(st.numpy() for st in states[i]) for i in range(len(states)) ] elif isinstance(states, list) and isinstance(states[0], list): states_np = [[ tuple(st.numpy() for st in states) for states in states[layer] ] for layer in range(num_layers)] else: assert False if isinstance(states_np, list): params[states_name] = convert_list_to_vmobj(states_np) else: params[states_name] = states_np run_and_compare(mod, params, pt_result)
script_module = do_trace(model) script_out = script_module(inp) assert len(out[0]) > 0 and len(script_out[0]) > 0 # compare bbox coord print(np.max(np.abs(out[0].numpy() - script_out[0].numpy()))) torch._C._jit_pass_inline(script_module.graph) torch.jit.save(script_module, name + ".pt") def convert_roi_align(): def _impl(inputs, input_types): spatial_scale = inputs[2] pooled_size = (inputs[3], inputs[4]) sampling_ratio = inputs[5] return relay.op.vision.roi_align(inputs[0], inputs[1], pooled_size, spatial_scale, sampling_ratio) return _impl script_module = torch.jit.load("mask_rcnn.pt") input_shapes = [("input", (1, 3, 300, 300))] custom_map = {'torchvision::roi_align': convert_roi_align()} from_pytorch(script_module, input_shapes, custom_map) """ NotImplementedError: The following operators are not implemented: ['aten::expand_as', 'aten::__and__', 'aten::meshgrid', 'aten::__interpolate', 'aten::scatter', 'aten::nonzero', 'aten::split_with_sizes', 'aten::log2', 'prim::ImplicitTensorToNum', 'aten::index', 'aten::_shape_as_tensor', 'aten::scalar_tensor'] """
models = [ SimpleIf(10, 20).eval(), NestedIf(10, 20).eval(), ScalarLoop().eval(), SimpleLoop().eval(), LoopWithIf().eval(), SimpleScalarWhileLoop().eval(), SimpleWhileLoop().eval(), NestedLoop().eval() ] for raw_model in models: script_module = torch.jit.script(raw_model) input_name = get_graph_input_names(script_module)[0] input_shapes = {input_name: (10, 20)} mod, params = from_pytorch(script_module, input_shapes) print(mod["main"]) executor = relay.create_executor("vm", mod=mod, ctx=tvm.cpu(0), target="llvm") evaluator = executor.evaluate() for i in range(5): inp = torch.rand(input_shapes[input_name], dtype=torch.float) with torch.no_grad(): pt_result = raw_model(inp.clone()) params[input_name] = inp.numpy() op_res = evaluator(**params) if not isinstance(pt_result, torch.Tensor): tvm_res = np.asscalar(op_res.asnumpy())