Any demo converting a pytorch model to relax with dynamic shape on GPU?

Development unity
1

As it says, I want to try out the dynamic shape support by converting a pytorch model into relax.

Is there a demo for GPU? I didn’t see one.

The mlc-llm project constructs a relax model, I can’t see how it adapts for dynamic shape.

Any demo for a small model? Thank you!

1 Like
2

@freshbird2023 were you able to find this?

Thanks!

3

There is a very simple dynamo example to convert torch dynamic model.

import torch
from tvm.relax.frontend.torch import relax_dynamo


class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(100, 10)

    def forward(self, x):
        return torch.nn.functional.relu(self.lin(x))


mod = MyModule()

opt_mod = torch.compile(mod, backend=relax_dynamo(), dynamic=True)

inp = torch.randn(10, 100)
with torch.no_grad():
    tvm_out = opt_mod(inp)

print(tvm_out)
torch.testing.assert_close(tvm_out, mod(inp))
1 Like
4

I’m also very interested in a relax example on GPU. Any demo would be very helpful. Thank you.

5

Hi, thank you for the answer ! I’ve seen the mlc-llm project and know some usage about relax.