Skip to content
This repository has been archived by the owner on May 22, 2023. It is now read-only.

[Relax][frontend] torch fx importer #419

Merged
merged 4 commits into from
Feb 10, 2023

Conversation

spectrometerHBH
Copy link
Contributor

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 [email protected]

spectrometerHBH and others added 3 commits February 9, 2023 13:18
Implements the Relax importer from PyTorch, using torch FX.

An example use of the importer is:

```python
from tvm.relax.frontend import from_pytorch

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 = {"input_1": ((128, 10), "float32")}

mod: tvm.IRModule = from_pytorch(torch_model, input_info)

```

---------

Co-authored-by: Ruihang Lai <[email protected]>
Copy link
Collaborator

@sunggg sunggg left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for bringing TorchFx frontend, @spectrometerHBH @MasterJH5574!
This is a great move forward for our next steps.
I reviewed in mlc repo so only a minor comment here.

tvm.ir.assert_structural_equal(mod, expected)


@tvm.testing.requires_gpu
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need gpu?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The CPU docker image doesn’t seem to have PyTorch installed, at least in the mlc repo.

@tqchen tqchen merged commit 923a214 into tlc-pack:relax Feb 10, 2023
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants