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