I try to print the script code of a IRModule, but it looks strange that I get the right result using mod.script but got None using mod.script(), why is that?
I see the definition of IRModule in python/tvm/ir/module.py, it defines script method, but I can’t find a member named script, but why calling mod.script worked?
import torch
import onnx
from tvm import relay
from tvm.relay.frontend.onnx import from_onnx
class DynamicNet(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, y, z):
tmp_a = torch.add(x, y)
tmp_b = torch.add(y, z)
return torch.add(tmp_a, tmp_b)
def torch_export():
model = DynamicNet()
dummy_input = tuple([torch.ones((1, 2))] * 3)
print(dummy_input)
input_names = ['x', 'y', 'z']
output_names = ['result']
torch.onnx.export(model, dummy_input, "multi_add.onnx", verbose=False, input_names=input_names, output_names=output_names)
def compile_onnx():
model = onnx.load('multi_add.onnx')
mod, params = from_onnx(model)
print(mod.script)
print(mod.script())
torch_export()
compile_onnx()
It shows as follows:
mod.script
<bound method IRModule.script of #[version = "0.0.5"]
def @main(%x: Tensor[(1, 2), float32] /* ty=Tensor[(1, 2), float32] */, %y: Tensor[(1, 2), float32] /* ty=Tensor[(1, 2), float32] */, %z: Tensor[(1, 2), float32] /* ty=Tensor[(1, 2), float32] */) -> Tensor[(1, 2), float32] {
%0 = add(%x, %y) /* ty=Tensor[(1, 2), float32] */;
%1 = add(%y, %z) /* ty=Tensor[(1, 2), float32] */;
add(%0, %1) /* ty=Tensor[(1, 2), float32] */
}
>
mod.script()
# from tvm.script import tir as T
@tvm.script.ir_module
class Module: