How to train GCN with a giant graph by torch_geometeric?

Question

I met an OOM problem when I was training a GCN layer with a graph where there are millions of nodes. Obviously, we cannot build an adjacency matrix of the graph to perfome message passing since the adjacency matrix, which contains terabytes, is too big. But in theory, we do can train it with much less memory by following steps:

1. for each node, perfome feature transformation.
2. for each node, aggregate features from its neighbors.

We can assume that matrix of all nodes' features can be stored in memory

pseudo code:

class GCNLayer(torch.nn.Module):
    def forward(self, x, edge_index):
        x = self.feature_transformation(x)
        x = self.aggregate(x, edge_index)
        return x

    def feature_transformation(self, x):
        return self.fc(x)

    def aggregate(self, x, edge_index):
        res = []
        for index, node in enumerate(x):
            adjacency = self.create_normolized_adjacency_vector_for(index, edge_index)
            res.append(adjacency * node)
        return torch.tensor(res, dtype=torch.float)

where create_normolized_adjacency_vector_for function is used to create a tensor of the index line of

But I don't know how to implement create_normolized_adjacency_vector_for method with high efficiency. Also, for loop seems to be slow too.

I wondering whether there is an efficient implementation, which is also easy to use in torch_geometric.