nn.TransformerConv leading CUDA out of memory?

[ShengzheXu]

I'm just running a simple test on ogbl-collab dataset. It seems when the node=235868 and edge=2358104, the following line raise a CUDA out of memory error. The code runs smoothly on small-scale demo input data.

It looks it's a multiplication of [2358104, 1, 256] * [2358104, 1, 1]? I'm new to the geometric toolbox; this is calculating the whole set of edges in the graph right? Appreciate for any helps.

(

pytorch_geometric/torch_geometric/nn/conv/transformer_conv.py

Line 219 in 129d60f

out *= alpha.view(-1, self.heads, 1)

)

[rusty1s]

TransformerConv scales linearly with the number of edges in your graph, so it is expected that there is a certain threshold to fit everything into GPU memory. I think you will need to use some scalability techniques to make TransformerConv applicable on larger graphs, e.g, via loader.NeighborLoader, loader.ClusterLoader or via loader.GraphSAINTSampler.

[wsad1]

Thanks, @rusty1s that's certainly the right thing to do for large graphs.
On another note I think moving the query, key and value computation outside of the message function might give some computation and memory gains for TransformerConv. I'll try to create a PR in a few days.

[rusty1s]

Yes, you are right. We omit the creation of x_i and x_j this way.