How to create torchscript models for GNNs?

[gunabam]

I tried implementing the following code but I am encountering the following error. Does anyone know how to solve this error? Is there a jittable function for BasicGNN class in pytorch geometric

torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:

import torch
from torch import nn
from torch_geometric.nn import GAT
from torch_geometric.data import Data

gnn = GAT(in_channels=256,
          hidden_channels=256,
          num_layers=4,
          heads=4,
          dropout=0.1,
          concat=False,
          v2=True,
          edge_dim=10,
          norm=nn.BatchNorm1d(256))

compiled_model = torch.jit.script(gnn)
torch.jit.save(compiled_model, 'compiled_model.pt')

I also tried something simpler:

import torch
from torch_geometric.nn import GCNConv
conv1 = GCNConv(16, 16).jittable()
print(torch.jit.script(conv1).graph)

But this also seems to be presenting the same error:

torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:

        def __init__(self, *args, **kwargs):
                                   ~~~~~~~ <--- HERE
            raise ImportError("'SparseTensor' requires 'torch-sparse'")

This is the version of my packages:

torch                     2.0.0                    pypi_0    pypi
torch-geometric           2.3.0                    pypi_0    pypi

[rusty1s]

For TorchScript support, you currently need to have torch-sparse installed. This will be fixed in PyG 2.4 :(

[rusty1s]

What does

import torch_geometric
print(torch_geometric.typing.WITH_TORCH_SPARSE)

return?

[rusty1s]

Also, this should be fixed in master already. If that's an option, just run

pip install git+https://github.com/pyg-team/pytorch_geometric.git

[gunabam]

Hi @rusty1s ,

So, when I run the following code, I am getting False.

import torch_geometric
print(torch_geometric.typing.WITH_TORCH_SPARSE)

I also built a new conda env with the following commands:

conda create --name pyG python==3.9
pip install git+https://github.com/pyg-team/pytorch_geometric.git
pip3 install torch torchvision torchaudio

I ran the code again with the new environment and I am getting False.

However, with the new conda env, when I run:

import torch
from torch_geometric.nn import GCNConv
conv1 = GCNConv(16, 16).jittable()
print(torch.jit.script(conv1).graph)

I get a new error:

    return self._c._get_method("forward").graph
RuntimeError: Method 'forward' is not defined.

[rusty1s]

That is expected since the input types are not fully resolved yet (GCNConv can operate both on edge_index and SparseTensor). It is fixable by either:

1. specifying the input types as additional argument to jittable:

import torch

from torch_geometric.nn import GCNConv

conv1 = GCNConv(16, 16)
conv1 = conv1.jittable('(Tensor, Tensor, OptTensor) -> Tensor')
print(torch.jit.script(conv1).graph)

2. Wrapping your layer inside a module with clearly specified types:

class GCN(torch.nn.Module):
    def __init__(self, ...):
        super().__init__()
        self.conv1 = GCNConv(16, 16)
        
    def forward(self, x, edge_index):
        return self.conv1(x, edge_index)