Is possible export the model on cpu and run it gpu?

torch.compile
1

As the title said. If no, how to export a model that can be runed on GPU.

import torch
from torch.export import export 

class ToyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(10, 10)
        self.fc2 = torch.nn.Linear(10, 10)
        self.fc3 = torch.nn.Linear(10, 10)
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        
        a = torch.sin(x)
        b = torch.cos(y)
        y = self.fc1(x)
        return a + b + y
f = ToyModule()

example_args = (torch.randn(10, 10), torch.randn(10, 10))
out_ref = f(*example_args)

exported_program: torch.export.ExportedProgram = export(
    f, args=example_args
)
print(exported_program) 


out = exported_program.module()(*example_args)
assert torch.allclose(out, out_ref)
2

Your code fails with:

ValueError: Expected `mod` to be an instance of `torch.nn.Module`, got <class 'function'>.

for me. Which PyTorch version are you using?

3

Oh, the export removed the support for function, and requires a module. [export] require Module to be passed to export by suo · Pull Request #117528 · pytorch/pytorch · GitHub

I revised my sample code, please try again. Thx!

My torch version is 2.3.0a0+git71d0202.

4

Thanks for updating the code!
Based on the new code it should work fine exporting the model on the CPU and loading and running it on the GPU:

torch.export.save(exported_program, "test.pt")
m = torch.export.load("test.pt")
print(m)

m.module()(*[e.cuda() for e in example_args])
# RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

m.module().cuda()
out = m.module()(*[e.cuda() for e in example_args])
assert torch.allclose(out.cpu(), out_ref)
1 Like
5

Cool, it works, thanks a lot!