tlc-pack / relax

Apache License 2.0
193 stars 58 forks source link

[Relax][frontend] torch fx importer #419

Closed spectrometerHBH closed 1 year ago

spectrometerHBH commented 1 year ago

Implements the Relax importer from PyTorch, using torch FX.

An example use of the importer is:

from torch import fx
from tvm.relax.frontend.torch import from_fx

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(in_features=10, out_features=7, bias=True)

    def forward(self, input):
        return self.linear(input)

torch_model = MyModule()
input_info = [((128, 10), "float32")]
gm = fx.symbolic_trace(torch_model)

mod: tvm.IRModule = from_fx(gm, input_info)

Co-authored-by: Ruihang Lai ruihangl@cs.cmu.edu