def torch2tvm_module(torch_module: torch.nn.Module,
torch_inputs: Tuple[torch.Tensor, ...], target):
torch_module.eval()
input_names = []
input_shapes = {}
with torch.no_grad():
for index, torch_input in enumerate(torch_inputs):
name = "i" + str(index)
input_names.append(name)
input_shapes[name] = torch_input.shape
buffer = io.BytesIO()
torch.onnx.export(
torch_module,
torch_inputs,
buffer,
input_names=input_names,
output_names=["o" + str(i) for i in range(len(torch_inputs))])
outs = torch_module(*torch_inputs)
buffer.seek(0, 0)
onnx_model = onnx.load_model(buffer)
relay_module, params = tvm.relay.frontend.from_onnx(onnx_model,
shape=input_shapes)
with tvm.relay.build_config(opt_level=3):
graph, tvm_module, params = tvm.relay.build(relay_module,
target,
params=params)
return graph, tvm_module, params
def from_torch(module, dummy_inputs):
if isinstance(dummy_inputs, torch.Tensor):
dummy_inputs = (dummy_inputs,)
input_shape = {}
for index, dummy_input in enumerate(dummy_inputs):
if isinstance(dummy_input, np.ndarray):
dummy_input = torch.from_numpy(dummy_input)
input_shape[str(index)] = dummy_input.shape
buffer = io.BytesIO()
module.eval()
torch.onnx.export(module, dummy_inputs, buffer)
buffer.seek(0, 0)
onnx_model = onnx.load_model(buffer)
return tvm.relay.frontend.from_onnx(onnx_model, shape=input_shape)
def torch2tvm_module(torch_module: torch.nn.Module,
torch_inputs: Tuple[torch.Tensor, ...], target):
torch_module.eval()
input_names = []
input_shapes = {}
with torch.no_grad():
for index, torch_input in enumerate(torch_inputs):
name = "i" + str(index)
input_names.append(name)
input_shapes[name] = torch_input.shape
buffer = io.BytesIO()
torch.onnx.export(
torch_module,
torch_inputs,
buffer,
input_names=input_names,
output_names=["o" + str(i) for i in range(len(torch_inputs))])
outs = torch_module(*torch_inputs)
buffer.seek(0, 0)
onnx_model = onnx.load_model(buffer)
relay_module, params = tvm.relay.frontend.from_onnx(onnx_model,
shape=input_shapes)
if AUTOTUNE:
graph, tvm_module, params = tr.tune_and_evaluate(
relay_module, params, target, torch_inputs)
elif USE_HISTORY:
with autotvm.apply_history_best('tsm1_save.log'):
with tvm.relay.build_config(opt_level=3):
graph, tvm_module, params = tvm.relay.build(relay_module,
target,
params=params)
else:
with tvm.relay.build_config(opt_level=3):
graph, tvm_module, params = tvm.relay.build(relay_module,
target,
params=params)
return graph, tvm_module, params
def torch2executor(torch_module: torch.nn.Module,
torch_inputs: Tuple[torch.Tensor, ...], target):
torch_module.eval()
input_names = []
input_shapes = {}
with torch.no_grad():
for index, torch_input in enumerate(torch_inputs):
name = str(index)
input_names.append(name)
input_shapes[name] = torch_input.shape
buffer = io.BytesIO()
torch.onnx.export(torch_module,
torch_inputs,
buffer,
input_names=input_names)
buffer.seek(0, 0)
onnx_model = onnx.load_model(buffer)
relay_module, params = tvm.relay.frontend.from_onnx(onnx_model,
shape=input_shapes)
with tvm.relay.build_config(opt_level=3):
graph, tvm_module, params = tvm.relay.build(relay_module,
target,
params=params)
ctx = tvm.gpu() if target.startswith('cuda') else tvm.cpu()
graph_module = graph_runtime.create(graph, tvm_module, ctx)
for pname, pvalue in params.items():
graph_module.set_input(pname, pvalue)
def executor(inputs: Tuple[np.ndarray]):
for index, value in enumerate(inputs):
graph_module.set_input(index, tvm.nd.array(value, ctx))
graph_module.run()
return tuple(
graph_module.get_output(index) for index in range(len(inputs)))
return executor