Loading and saving TorchScript message passing networks

[dav0dea]

Hi,
is it possible to load and save jittable message passing modules through torch.save and torch.load?

Example:

import torch
from torch_geometric.nn import GCNConv

layer = GCNConv(64, 64).jittable()
torch.save(layer, "conv.pt")

Then, when I run torch.load("conv.pt") I get ModuleNotFoundError: No module named 'GCNConvJittable_c3e7ca'. Interestingly enough, when I load the file from the same process immediately after saving it works.

My current setup is to save all arguments required for constructing the model together with the model's state dict and reconstructing a new model from the arguments. I feel like it would be more elegant and intuitive to just use the load and save functions as provided by PyTorch.
Do you have a solution for this or are you planning to implement this in the future?

[rusty1s]

I see. Thanks for letting me know. This is indeed a current limitation as we are saving a temporary file of the jittable module to disk and load it into Python afterwards. I will look into ways how to circumvent this, but for now, I suggest to use regular load_state_dict to initialize the model.

[dav0dea]

Alright, thank you!

[raimis]

As far as I check, if you convert the layer to TorchScript, you can save and load it without that error:

import torch
from torch_geometric.nn import GCNConv

layer = GCNConv(64, 64).jittable()
torch.jit.script(layer).save('conv.pt')

new_layer = torch.jit.load('conv.pt')
``

[raimis]

@11carlesnavarro this might be a solution to your problem.