NotSupportedError when using jittable().

[sunwoongc]

I'm trying to use PyG with Pytorch-lightning and deploy the model with BentoML. During saving model, I got some error related to torchscript(script_module = model.to_torchscript()).

First I'll simply my problem here.

import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GNN(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, 64).jittable()
        self.conv2 = GCNConv(64, out_channels).jittable()

    def forward(self, x, *args):
        edge_index = args[0]
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

In the above code, I got the error:

NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "<ipython-input-17-34dbcc76746f>", line 10
    def forward(self, x, *args):
                         ~~~~~ <--- HERE
        edge_index = args[0]
        x = self.conv1(x, edge_index)

First, can't I use args, kwargs?

---

And this is part of my implementation.

class GNNModels(torch.nn.Module):
    def __init__(self, conv=None, norm=None, act=None, block='res+', 
                 dropout=0., ckpt_grad=False,edge_norm = None,):
        super(DeepGCNLayerV2, self).__init__()

        self.conv = conv.jittable()
        self.norm = norm
        self.act = act
        self.edge_norm  = edge_norm
        self.block = block.lower()
        assert self.block in ['res+', 'res', 'dense', 'plain']
        self.dropout = dropout
        self.ckpt_grad = ckpt_grad


    def reset_parameters(self):
        self.conv.reset_parameters()
        self.norm.reset_parameters()

    def forward(self, x, edge_index, edge_attr):
        """"""
        if isinstance(self.conv, SAGEConvV2): # Follow the original implementation.
            org_edge_attr = edge_attr #org_edge_attr:[edge_index,edge_attr,....]
            if org_edge_attr is None: org_edge_attr = 0
            if self.block == 'res+':
                if self.norm is not None:

After I got the same error as above, I removed args and kwargs, and using jittable() in __init__. And then I got the error:

RuntimeError: 
'__torch__.SAGEConvV2Jittable_b8fd43.Propagate_b8fd43' object has no attribute or method 'edge_attr'.:
  File "/tmp/tony_pyg/tmpyl682gnd.py", line 126
            self.__set_size__(size, 1, x_i)
            x_i = self.__lift__(x_i, edge_def, i)
        edge_attr = kwargs.edge_attr
                    ~~~~~~~~~~~~~~~~ <--- HERE
        x_j: torch.Tensor = init
        data = kwargs.x

and SAGEConvV2 is

class SAGEConvV2(MessagePassing):
    # concat edge_index and node features
    def __init__(
        self,
        in_channels: Union[int, Tuple[int, int]],
        out_channels: int, 
        normalize: bool = False,
        root_weight: bool = True,
        bias: bool = True,
        in_edge_channels: Union[int, Tuple[int, int],None] = None ,
        aggr: str = 'mean', 
        t: float = 1.0, 
        learn_t: bool = False,
        p: float = 1.0, 
        learn_p: bool = False,
        **kwargs
    ): 
...

How can I handle this issue?

(Full code is in this repository, I modify a little bit of a code, but not that much.)

[rusty1s]

Yes, you cannot use a variable number of args since PyTorch will fail to compile this.

For your own MessagePassing Instantiation to be jittable, there are also a few things you need to take care of. They are described here.

[sunwoongc]

Thanks for the help! It works for me, too 😄.
Then I got and error,

torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/home/tony/anaconda3/envs/deepdrug-deploy/lib/python3.8/site-packages/torch/utils/checkpoint.py", line 164
def checkpoint(function, *args, use_reentrant: bool = True, **kwargs):
                                                             ~~~~~~~ <--- HERE
    r"""Checkpoint a model or part of the mode

But it seems that it's a problem from PyTorch's checkpoint.py. Are there any possibilities that the problem come from PyG?

[rusty1s]

No, this comes from PyG since TorchScript always needs clearly defined types at the top-level, but there are two Union there: x: Union[Tensor, PairTensor] and edge_index: Union[Tensor, SparseTensor]. You should be able to resolve this by just doing

def forward(self, x: Tensor, edge_index: Tensor, edge_attr: Optional[Tensor] = None, size: Size = None) -> Tuple[Tensor, Tensor]:

[sunwoongc]

I can't find any edge_index: Union[Tensor, SparseTensor]. 😢

class SAGEConvV2(MessagePassing):
    # concat edge_index and node features
    def __init__(
        self,
        in_channels: Union[int, Tuple[int, int]],
        out_channels: int,
        normalize: bool = False,
        root_weight: bool = True,
        bias: bool = True,
        in_edge_channels: Union[int, Tuple[int, int], None] = None,
        aggr: str = 'mean',
        t: float = 1.0,
        learn_t: bool = False,
        p: float = 1.0,
        learn_p: bool = False,
        **kwargs
    ):  # yapf: disable
        # kwargs.setdefault('aggr', 'mean')

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.normalize = normalize
        self.root_weight = root_weight
        self.in_edge_channels = in_edge_channels

        if isinstance(in_channels, int):
            in_channels = (in_channels, in_channels)

        kwargs.setdefault('aggr', None)
        super().__init__(**kwargs)

        if self.root_weight:
            self.lin_r = Linear(in_channels[1], out_channels, bias=False)

        self.aggr = aggr
        self.eps = 1e-7
        assert aggr in ['softmax', 'softmax_sg', 'power', 'mean', 'add', 'sum']
        self.t = None
        self.p = None
        if self.aggr in [
                'softmax',
                'softmax_sg',
                'power',
        ]:
            self.initial_t = t
            if learn_t and aggr == 'softmax':
                self.t = Parameter(torch.Tensor([t]), requires_grad=True)
            else:
                self.t = t
            self.initial_p = p
            if learn_p:
                self.p = Parameter(torch.Tensor([p]), requires_grad=True)
            else:
                self.p = p

        if in_edge_channels is not None:
            self.lin_l = nn.Sequential(
                nn.Linear(in_channels[0] + in_edge_channels,
                          in_channels[0] * 2, bias=bias),
                nn.ReLU(),
                nn.Linear(in_channels[0] * 2, out_channels, bias=bias),
            )
        else:
            self.lin_l = nn.Sequential(
                nn.Linear(in_channels[0], in_channels[0] * 2, bias=bias),
                nn.ReLU(),
                nn.Linear(in_channels[0] * 2, out_channels, bias=bias),
            )

    def message(self, x_i: Tensor, x_j: Tensor, edge_attr: Tensor) -> Tensor:
        if edge_attr is None:
            x = x_j
        else:
            x = torch.cat([x_j, edge_attr], dim=-1)
        x = self.lin_l(x)
        return x

    def aggregate(self, inputs: Tensor, index: Tensor,
                  dim_size: Optional[int] = None) -> Tuple[Tensor, Tensor]:

        if self.aggr == 'softmax':
            assert self.t is not None
            out = scatter_softmax(inputs * self.t, index, dim=self.node_dim)
            return scatter(inputs * out, index, dim=self.node_dim,
                           dim_size=dim_size, reduce='sum'), inputs

        elif self.aggr == 'softmax_sg':
            assert self.t is not None
            out = scatter_softmax(inputs * self.t, index,
                                  dim=self.node_dim).detach()
            return scatter(inputs * out, index, dim=self.node_dim,
                           dim_size=dim_size, reduce='sum'), inputs

        elif self.aggr == 'power':
            assert self.p is not None
            min_value, max_value = 1e-7, 1e1
            torch.clamp_(inputs, min_value, max_value)
            out = scatter(torch.pow(inputs, self.p), index, dim=self.node_dim,
                          dim_size=dim_size, reduce='mean')
            torch.clamp_(out, min_value, max_value)
            return torch.pow(out, 1 / self.p), inputs

        else:
            return scatter(inputs, index, dim=self.node_dim, dim_size=dim_size,
                           reduce=self.aggr), inputs

    def update(self, inputs: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
        return inputs

    def forward(self, x: Tensor, edge_index: Tensor,
                edge_attr: Optional[Tensor] = None,
                size: Size = None) -> Tuple[Tensor, Tensor]:
        """"""
        if isinstance(x, Tensor):
            x: PairTensor = (x, x)
        # propagate_type: (x: PairTensor, edge_attr: OptTensor)
        out, edge_attr = self.propagate(edge_index, x=x, edge_attr=edge_attr,
                                        size=size)

        x_r = x[1]
        if self.root_weight and x_r is not None:
            out += self.lin_r(x_r)

        if self.normalize:
            out = F.normalize(out, p=2., dim=-1)

        return out, edge_attr

Here's my current SAGEConvV2.

Thanks for the help.

[rusty1s]

Yes, this looks correct to me. Does it work?

BTW: Adj is defined as Union[Tensor, SparseTensor in torch_geometric/typing.py.

[sunwoongc]

Hello! I checked that the SAGEConv class works fine. But I got the same error,

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

def checkpoint(function, *args, use_reentrant: bool = True, **kwargs):
                                                             ~~~~~~~ <--- HERE
    r"""Checkpoint a model or part of the model

but it happened inside of class DeepGCNLayerV2:

class DeepGCNLayerV2(torch.nn.Module):
    def __init__(self, 
                conv: Optional[Module] = None,
                norm: Optional[Module] = None, 
                act: Optional[Module] = None,
                edge_norm: Optional[Module] = None,
                block: str = 'res+', 
                dropout: float = 0.,
                ckpt_grad: bool = False):

        super(DeepGCNLayerV2, self).__init__()

        self.conv = conv.jittable() # NOTE: Change for TorchScript
        self.norm = norm
        self.act = act
        self.edge_norm  = edge_norm
        self.block = block.lower()
        assert self.block in ['res+', 'res', 'dense', 'plain']
        self.dropout = dropout
        self.ckpt_grad = ckpt_grad


    def reset_parameters(self):
        self.conv.reset_parameters()
        self.norm.reset_parameters()

    def forward(self, x: Tensor, edge_index: Tensor, edge_attr: Tensor) -> Tensor:
        """"""
        # args = list(args)
        # x = args.pop(0)
        # if isinstance(self.conv, SAGEConvV2): # Follow the original implementation.

        org_edge_attr = edge_attr #org_edge_attr:[edge_index,edge_attr,....]
        if org_edge_attr is None: org_edge_attr = torch.tensor(0)
        
        if self.block == 'res+':
            h = x
            if self.norm is not None:
                h = self.norm(h)
            if self.act is not None:
                h = self.act(h)
            if self.edge_norm is not None:
                edge_attr = self.edge_norm(edge_attr)
            h = F.dropout(h, p=self.dropout, training=self.training)
            if self.conv is not None and self.ckpt_grad and h.requires_grad:
                h,edge_attr = checkpoint(self.conv, h, edge_index, edge_attr)
            else:
                h,edge_attr = self.conv(h, edge_index, edge_attr)

            return x+h, org_edge_attr+edge_attr

        else:
            if self.conv is not None and self.ckpt_grad and x.requires_grad:
                h = checkpoint(self.conv, x, edge_index, edge_attr)
            else:
                h = self.conv(x, edge_index, edge_attr)
            if self.norm is not None:
                h = self.norm(h)
            if self.act is not None:
                h = self.act(h)

            if self.block == 'res':
                return x + h
            elif self.block == 'dense':
                return torch.cat([x, h], dim=-1)
            elif self.block == 'plain':
                pass

            return F.dropout(h, p=self.dropout, training=self.training)

    def __repr__(self):
        return '{}(block={})'.format(self.__class__.__name__, self.block)

Here, I changed a code little bit, I modified self.conv = conv.jittable(), in this case conv is SAGEConvV2. And torch.jit.script(self.conv) works fine. I try to add TypeHint in the forward method of DeepGCNLayerV2 as you mentioned before. But I still can't fix the error.

I used DeepGCNLayerV2 in another class called DeepGCN as follows:

...
 for block in range(num_blocks):
            layers = nn.ModuleList()
            for i in range(1, num_layers + 1):
                conv = graph_conv(mid_channel, mid_channel, **graph_para_dict)
                norm = nn.LayerNorm(mid_channel, elementwise_affine=True)
                edge_norm = nn.LayerNorm(mid_channel, elementwise_affine=True)
                act = nn.ReLU(inplace=True)

                layer = DeepGCNLayerV2(conv, norm, act, block='res+', dropout=dropout_ratio,
                                     # ckpt_grad=  False,
                                     ckpt_grad=i % 3, edge_norm=edge_norm,
                                     )
                layers.append(layer)
...

here by debugging, I tried torch.jit.script(layer), it returns the same error(NotSupportedError in torch.utils.checkpoint.checkpoint).