Using edge attr

[JindeShubhamA]

Hello,

I have spent around 2 months on this repository and done some benchmarking on synthetic datasets for both Graph and Node classification tasks using GCN, SageConv, and Attention-based Graph networks. I would say this repository is the best that is out there for performing experiments on Graph-based networks.

The dataset I worked on till now did not have edge attributes. However know as I have become somewhat comfortable, I would like to also experiment with the edge attribute task. Edge attributes will be the node features that are passed during message propagation.

1. Can you direct on some examples where edge attributes are used?
2. Also are there some GNN layers that use edge attributes in their default implementation?

[rusty1s]

Thanks for your kind words :)

Not all GNN operators support edge features. This is mostly due to following the precise formulation of a given paper. While on can naturally incorporate edge features in the message passing phase, there exist multiple ways to do so (e.g. via summation or concatenation). The best way to find all GNN operators that can make use of edge features is to search for edge_attr in the torch_geometric.nn documentation. Here are a few out-of-my-head: GENConv, GINEConv, GMMConv, TransformerConv, PNAConv, ...

Here is an example that makes use of edge features.

[rusty1s]

Currently yes, you need to create a fully new-version of it. It might be also possible to implement something like this by making use of internal state, i.e.:

class MyGCNConv(GCNConv):
    def _init__(self, *args, edge_dim=0, **kwargs):
        super().__init__(*args, **kwargs):
        self.edge_lin = Linear(edge_dim, self.out_channels)
        
    def forward(self, x, edge_index, edge_attr):
        self._edge_attr = edge_attr
        out = super().forward(x, edge_index)
        self._edge_attr = None
        return out

    def message(self, x_j, edge_weight):
        return (x_j + self.edge_lin(self._edge_attr))) * edge_weight.view(-1, 1)

I need to think of better ways to handle such customization though.

[JindeShubhamA]

Found a way around. It is getting trained. Only issue is GCNConv adds self loop which increases number of edges and edge_attr does not have any values for it. Currently training by removing the self loop which does not make much sense. But yes I can tackle this through multiple ways. Thanks

[a-ism256]

Hello,

Do you mind sharing what ended up working for you?

I'm trying to do the same thing for SAGEConv and this thread seems to be the closest to verifying a working solution.

According to what @rusty1s said above, I'm implementing as follows. Also, I understand that this is the summation approach not the concatenation. Not sure if concatenation would be worth trying.

class EdgeSAGEConv(SAGEConv):
    def _init__(self, *args, edge_dim=0, **kwargs):
        super().__init__(*args, **kwargs)
        self.edge_lin = Linear(edge_dim, self.out_channels)
        
    def forward(self, x, edge_index, edge_attr):
        self._edge_attr = edge_attr
        out = super().forward(x, edge_index)
        self._edge_attr = None
        return out

    def message(self, x_j, edge_weight):
        return (x_j + self.edge_lin(self._edge_attr)) * edge_weight.view(-1, 1)

Then, the rest of the model initialization and training would be as in the reddit example. I haven't tested it yet, just trying to understand everything and make sure it all makes sense before I do.

Thanks!

[rusty1s]

Yes, this should work :)

[JindeShubhamA]

Hello,
I did try the method you are trying but had few issues. So I had a new version of GCNConv in which I am taking care of edge_attr as node features. As I wanted the same for SAGEConv, GATConv and TAGConv, I ended up spending sometime looking at the implementation of this standard layers and just modified them.