Use of nn.models.GIN for Graph Classification

[simon-forb]

It is great to see that the GIN model is implemented natively in PyTorch Geometric.

As I understand it however, this implementation cannot be used for graph classification as described in the original publication [1] because that would require the output of all GNN layers including the input layer (cf. Eq. 4.2). This is because GIN only yields the output of the last layer and the hidden representations are inaccessible.

Is that correct or am I missing something?

If this is correct, does that mean the best way to use GIN for graph classification is to re-implement the model using GINConv?

---

[1] Xu, Keyulu, et al. "How powerful are graph neural networks?." arXiv preprint arXiv:1810.00826 (2018). https://arxiv.org/abs/1810.00826

[rusty1s]

You can well use GIN for graph classification. Here is an example.

[simon-forb]

Thank you for your answer and the provided link which is useful to me. I already know from the original publication that GIN can indeed be used for graph classification (cf. Eq. 4.2).

However in my question, I was interested in whether there exists a way to use the class nn.models.GIN from PyTorch Geometric directly for this purpose. I thought that there was maybe another implementation of GIN available or that this was achievable through some class parameters which I missed.

As I understand your answer, this is not the case. Therefore one has to manually implement a custom class (GIN in your example) using the provided GINConv class.

Anyway, thanks again!

[rusty1s]

This would look as follows:

class Model(torch.nn.Module):
    def __init__():
         self.gnn = GIN(...)
         self.classifier = MLP(...)
         
     def forward(self, x, edge_index, batch):
          x = self.gnn(x, edge_index)
          x = global_mean_pool(x, batch)
          return self.classifier(x)