Memory efficient graph classification when graph is constant across examples

[mdanb]

I'm playing around with a dataset for graph classification where the graph is the same for all training instances, but the features are different from one training instance to the next. I'm trying to find a way to make Pytorch Geometric work for this case without duplicating the edges for each training instance so as to be as memory efficient as possible. How might I approach doing this?

[wsad1]

@mdanb, Iam trying to understand why you need to duplicate edges for each training instance, you could pass the same edge_index variable every time you call forward. Maybe you could paste a short snippet of your training code.

Also check out layers like GCNConv, if cached is set to true then the computation of normalizations is cached.

[mdanb]

See answer and comments below

[rusty1s]

Some GNN operators have static graph support built-in, e.g., in GCNConv you can do:

x = [batch_size, num_nodes, num_features]
edge_index = [2, num_edges]  # Edge-connectivity of a single graph

out = GCNConv(...)(x, edge_index)
# out will have shape [batch_size, num_nodes, num_new_features]

To be even more memory-efficient, I suggest you to use the SparseTensor class for encoding edges, see here.

[mdanb]

I see. So just to make sure I got this right, instead of doing something like (which I guess here is memory efficient but does not result in correct message passing):

x, batch, edge_index, edge_attr = data.x, data.batch, self.edge_index, self.edge_attr
for i in range(self.num_layers):
     x = self.convs[i](x, edge_index)

I would do this instead (just added a line to reshape x):

x, batch, edge_index, edge_attr = data.x, data.batch, self.edge_index, self.edge_attr
x = x.reshape(batch_size, num_nodes, num_features_per_node)
for i in range(self.num_layers):
     x = self.convs[i](x, edge_index)

[rusty1s]

Yes :)

[mdanb]

By the way, does GINE conv support this?

[rusty1s]

Yes, I added some tests to verify: https://github.com/rusty1s/pytorch_geometric/blob/master/test/nn/conv/test_gin_conv.py#L110-L128