Using edge features for GCN in DGL

I’m trying to implement a simple GCN. In many papers, edges have discrete features, and each possible value is associated with a different weight matrix or set of weight matrices. An example would be here. Is anyone familiar with how to implement a model like this in DGL? The DGL team’s example of GCNs for graph classification doesn’t use edge features, neither does another example I found online.

2 Likes

Recently we’ve released a model zoo on Chemistry and this might be related to your request: https://github.com/mufeili/dgl/tree/master/examples/pytorch/model_zoo/chem. In particular, the regression models in the model zoo have utilized edge features.

1 Like

Hey @mufeili,

How could I incorporate the edge features to be considered for node classification within the GCN class below?

class GCN(nn.Module):
  def __init__(self, in_feats, n_hidden, n_classes, n_layers, activation, dropout):
    super(GCN, self).__init__()
    self.layers = nn.ModuleList()
    # input layer
    self.layers.append(GraphConv(in_feats, n_hidden, activation=activation))
    # hidden layers
    for i in range(n_layers - 1):
        self.layers.append(GraphConv(n_hidden, n_hidden, activation=activation))
    # output layer
    self.layers.append(GraphConv(n_hidden, n_classes))
    self.dropout = nn.Dropout(p=dropout)

  def forward(self, g_dgl, features):
    h = features
    for i, layer in enumerate(self.layers):
        if i != 0:
            h = self.dropout(h)
        h = layer(g_dgl, h)
    return h

My original networks graph has three features for every nodes and I have converted them into DGL graph with this syntax
g_dgl.from_networkx(g_nx, edge_attrs=['featA','featB','featC'])

Thanks for the suggestion

1 Like

Depends on how your edge feature looks like.

If your edge feature is a single decimal number, one can treat the graph as a weighted graph. For GCNs, the messages sent from a node would be scaled by the squared root of weighted degree (i.e. the sum of incoming edge weights) of both endpoints. We are currently in the progress of adding this to various applicable modules.

If your edge feature is a single categorical variable, the model would be similar to RGCN. You can look at the example we provided.

Otherwise (e.g. your edge feature is a vector, or some combination of categorical and numeric variables) you will need to define the “message function” yourself. A simple option I would try is to define the message function as an MLP of source and edge features (and maybe destination features). An example would be:

class GNNLayer(nn.Module):
    def __init__(self, node_dims, edge_dims, output_dims):
        self.W_msg = nn.Linear(node_dims + edge_dims, output_dims)
        self.W_apply = nn.Linear(output_dims * 2, output_dims)
    def message_func(edges):
        return {'m': F.relu(self.W_msg(th.cat([edges.src['h'], edges.data['h']], 0)))}
    def forward(self, g, node_features, edge_features):
        with g.local_scope():
            g.ndata['h'] = node_features
            g.edata['h'] = edge_features
            g.update_all(self.message_func, fn.sum('m', 'h_neigh'))
            g.ndata['h'] = F.relu(self.W_apply(th.cat([g.ndata['h'], g.ndata['h_neigh']], 0)))
            return g.ndata['h']
3 Likes

Hey @BarclayII,

I tried the GNNLayer class to replace my GCN class without success. But I’m not sure where the mistake I made.

The size of my edge features' are:
featA: torch.Size([346320, 3])
featB: torch.Size([346320, 3])
featC: torch.Size([346320, 3])
The size of my node feature is: 
featNode: torch.Size([116054, 3, 3]) 
and The size of my output is: 
targetLabel: torch.Size([116054, 1])

I tried to create a model with model = GNNLayer(3, 3, 1), but it produces cannot assign module before Module.__init__() call.

Also, how do I know how many hidden layers I have when using GNNLayer class you created above? With the GCN class i mentioned above (from GCN example), we know how many hidden layers we have with for i in range(n_layers - 1): self.layers.append(GraphConv(n_hidden, n_hidden, activation=activation)).

Thanks so much for your help.

My bad. In the __init__ function you need to write super().__init__() to initialize the PyTorch NN Module. Also I didn’t test the module implementation; this is just giving you an idea of how to utilize the edge features.

Thanks @BarclayII,

I really appreciate your help. Indeed, it works with super().__init__(). As I’m still learning how to use this framework properly, there are many things that I’m not sure yet if my implementation even properly constructed. So far, what I have is as follow:

class GNNLayer(nn.Module):
  def __init__(self, node_dims, edge_dims, output_dims, activation):
    super(GNNLayer, self).__init__()
    self.W_msg = nn.Linear(node_dims + edge_dims, output_dims)
    self.W_apply = nn.Linear(output_dims * 2, output_dims)
    self.activation = activation  #I'm not sure if I need to put activation here??
  def message_func(edges):
    return {'m': F.relu(self.W_msg(th.cat([edges.src['h'], edges.data['h']], 0)))}
  def forward(self, g, node_features, edge_features):
    with g.local_scope():
      g.ndata['h'] = node_features
      g.edata['h'] = edge_features
      g.update_all(self.message_func, fn.sum('m', 'h_neigh'))
      g.ndata['h'] = F.relu(self.W_apply(th.cat([g.ndata['h'], g.ndata['h_neigh']], 0)))
      return g.ndata['h']

class GCN(nn.Module):
  def __init__(self, dim_in, dim_out, activation, dropout):
    super(GCN, self).__init__()
    self.layers = nn.ModuleList()
    #self.layers.append(GraphConv(22, 16, activation=activation)) #comment as example
    #self.layers.append(GraphConv(16, 5, activation=activation)) #comment as example
    self.layers.append(GNNLayer(3,3,1,activation=activation))
    self.layers.append(GNNLayer(3,3,1,activation=activation))
    self.dropout = nn.Dropout(p=dropout)
  def forward(self, g_dgl, features):
    h = features
    for i, layer in enumerate(self.layers):
        if i != 0:
            h = self.dropout(h)

As commented above:

  1. Is it correct to put activation in GNNLayer class, which I assume would be the same as when we use GraphConv() class when we add a new hidden layer?
  2. With self.layers.append(GraphConv(16, 5, activation=activation)), it is clear that the output_dim of previous layer (i.e.16) becomes the input_dim for the next layer. But when GNNLayers takes node and edge features, how is the same principle is applied (output_dim becomes input_dim)?

I guess for the implementation part, I will ask in the new thread with more details. Thanks so much for the help. I appreciate it.

  1. It seems that you do not use self.activation anywhere.
  2. This is answered here.

Thanks @mufeili,
I use self.activation in GNNLayer class in def init. Should I use it there? or in other function in GNNLayer class?

You only defined it in __init__, but you did not use it at all in forward. I’d recommend you going through some tutorials in PyTorch here.