Memory efficient GATConv

[HongtaoYang]

Hi, I've been using GAT on very dense graphs where number of edges (E) is very close to N^2 (N is number of nodes). The issue is that during the propagation process, x_i and x_j got materialised as tensors of shape [E, feature_dim]. This takes a lot of memory. I asked a related question #2356

I tried to work around this issue using SparseTensor and implements message_and_aggregate method. Here is the modified GATConv code:

class GATConv(MessagePassing):

    ...

    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                size: Size = None, return_attention_weights=None, **kwargs):
        # type: (Union[Tensor, OptPairTensor], Tensor, Size, NoneType) -> Tensor  # noqa
        # type: (Union[Tensor, OptPairTensor], SparseTensor, Size, NoneType) -> Tensor  # noqa
        # type: (Union[Tensor, OptPairTensor], Tensor, Size, bool) -> Tuple[Tensor, Tuple[Tensor, Tensor]]  # noqa
        # type: (Union[Tensor, OptPairTensor], SparseTensor, Size, bool) -> Tuple[Tensor, SparseTensor]  # noqa
        r""" Mostly the same as in the parent class, only difference is added skip connection and layer norm.

        Args:
            return_attention_weights (bool, optional): If set to :obj:`True`,
                will additionally return the tuple
                :obj:`(edge_index, attention_weights)`, holding the computed
                attention weights for each edge. (default: :obj:`None`)
        """
        H, C = self.heads, self.out_channels

        x_l: OptTensor = None
        x_r: OptTensor = None
        alpha_l: OptTensor = None
        alpha_r: OptTensor = None
        if isinstance(x, Tensor):
            assert x.dim() == 2, 'Static graphs not supported in `GATConv`.'
            x_l = x_r = self.lin_l(x).view(-1, H, C)
            alpha_l = (x_l * self.att_l).sum(dim=-1)
            alpha_r = (x_r * self.att_r).sum(dim=-1)
        else:
            x_l, x_r = x[0], x[1]
            assert x[0].dim() == 2, 'Static graphs not supported in `GATConv`.'
            x_l = self.lin_l(x_l).view(-1, H, C)
            alpha_l = (x_l * self.att_l).sum(dim=-1)
            if x_r is not None:
                x_r = self.lin_r(x_r).view(-1, H, C)
                alpha_r = (x_r * self.att_r).sum(dim=-1)

        assert x_l is not None
        assert alpha_l is not None

        num_nodes = x_l.size(0)
        if x_r is not None:
            num_nodes = min(num_nodes, x_r.size(0))
        if size is not None:
            num_nodes = min(size[0], size[1])
        
        if self.add_self_loops:
            if isinstance(edge_index, Tensor):
                edge_index, _ = remove_self_loops(edge_index)
                edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
            elif isinstance(edge_index, SparseTensor):
                edge_index = set_diag(edge_index)

        if isinstance(edge_index, SparseTensor):
            row, col, value = edge_index.coo()
            alpha = alpha_l[row] if alpha_r is None else alpha_l[row] + alpha_r[col]
            alpha = F.leaky_relu(alpha, self.negative_slope)
            alpha = softmax(alpha, col)
            alpha = F.dropout(alpha, p=self.dropout, training=self.training)
            alpha_all_heads = [SparseTensor(row=row, col=col, value=alpha[:, i], sparse_sizes=(num_nodes, num_nodes)).t() for i in range(H)]
            out = self.propagate(edge_index, x=x_l, alpha=alpha_all_heads, size=size)

        else:
            out = self.propagate(edge_index, x=(x_l, x_r), alpha=(alpha_l, alpha_r), size=size)

        ...



    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor, alpha: List[SparseTensor]):
        x_heads = [matmul(a, x[:, i, :], reduce=self.aggr) for i, a in enumerate(alpha)]
        x_new = torch.cat(x_heads, dim=-1)

        return x_new

Note that I still have to materialize the attention as tensor of shape [E, num_heads] because attentions are different for each edge. But node features don't need to be expanded for each edge.

Any thoughts on this? Thanks!

[rusty1s]

Yes, this works. Good job :) The downside of this approach is that we need to manually loop over each head.

[HongtaoYang]

Any ideas on how to improve regarding multi-heads? Does pytorch_sparse support higher dimensional matrix multiplication?

[rusty1s]

Not yet, this would be something we need to add to the pytorch_sparse package.

[HongtaoYang]

Thanks! With current way of handling multi-heads, do you think its worth creating a pr for?

[rusty1s]

Yes, please do so. We can create a MemoryEfficientGATConv, for example.

[thegodone]

nice, could be also extend to AttFP