torch_geometric.nn.models for classification

[Takaogahara]

Hello everyone,

I've been looking in the documentation and Github issues for some time now and haven't been able to find a way to use the models available in torch_geometric.nn such as GCN for graph classification.

Is there any way to use them or are they just for node classification?

Thanks!

[rusty1s]

You can use the models in nn.models as part of graph classification pipeline:

x = self.gcn(x, edge_index)
x = global_mean_pool(x, batch)
return self.classifier(x)

We assume that a GNN just learns new node representations, which can then be used for various down-stream tasks (node, edge, graph) by connecting different heads to it.

[Takaogahara]

Thanks!!!