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
@freshbird2023 were you able to find this?
Thanks!
lqf
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
Amy
4
I’m also very interested in a relax example on GPU. Any demo would be very helpful. Thank you.
Hi, thank you for the answer ! I’ve seen the mlc-llm project and know some usage about relax.