How to use mutiple node features and edge features in GCN?

[Silas]

Hello, I’m searching how to use GCN on graph classification with both node features and edge features. However I’m truly a rookie in GNNs, and I’ve been struggling in dealing with the codes.

I found a topic earlier here: Node and edge features in with GCN

import dgl
import dgl.function as fn
import torch
import torch.nn as nn
import torch.nn.functional as F

class GNNLayer2(nn.Module):
    def __init__(self, ndim_in, edims, ndim_out, activation):
        super(GNNLayer2, self).__init__()
        self.W_msg = nn.Linear(ndim_in + edims, ndim_out)
        self.W_apply = nn.Linear(ndim_in + ndim_out, ndim_out)
        self.activation = activation

    def message_func(self, edges):
        return {'m': F.relu(self.W_msg(torch.cat([edges.src['h'], edges.data['h']], 2)))}

    def forward(self, g_dgl, nfeats, efeats):
        with g_dgl.local_scope():
            g = g_dgl
            g.ndata['h'] = nfeats
            g.edata['h'] = efeats
            g.update_all(self.message_func, fn.sum('m', 'h_neigh'))
            g.ndata['h'] = F.relu(self.W_apply(torch.cat([g.ndata['h'], g.ndata['h_neigh']], 2)))
            return g.ndata['h']


class GCN2(nn.Module):
    def __init__(self, ndim_in, ndim_out, edim, activation, dropout):
        super(GCN2, self).__init__()
        self.layers = nn.ModuleList()
        self.layers.append(GNNLayer2(ndim_in, edim, 50, activation))
        self.layers.append(GNNLayer2(50, edim, 25, activation))
        self.layers.append(GNNLayer2(25, edim, ndim_out, activation))
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, g, nfeats, efeats):
        for i, layer in enumerate(self.layers):
            if i != 0:
                nfeats = self.dropout(nfeats)
            nfeats = layer(g, nfeats, efeats)
        return nfeats.sum(1)

if __name__ == '__main__':
    model = GCN2(3, 1, 3, F.relu, 0.5)
    g = dgl.DGLGraph([[0, 2], [2, 3]])
    nfeats = torch.randn((g.number_of_nodes(), 3, 3))
    efeats = torch.randn((g.number_of_edges(), 3, 3))
    model(g, nfeats, efeats)

I read this code, and my question is, what should I do if I have 2 kinds of node data, of which one is a 200d vector, and the other is an 8d vector. Besides, I also have an edge data, which is an 8d vector. How can I put them together when building GCN layer?

I’m sorry to ask these silly questions. I would appreciate if you could kindly help.

[mufeili]

1. I think you can simply concatenate the two kinds of node data.
2. The model in the code snippet handles both node and edge features. nfeats for node features and efeats for edge features.

[Silas]

Thank you for your reply.
Another questions

1. if I want to adjust the codes to do graph classification (n_classes=3), is it okay to code like this? (Because I want to use crossentropy loss with softmax so I just add a fully-connected layer on top of the original model)

    class GNNLayer(nn.Module):
    def __init__(self, ndim, edim, out_dim): 
        super(GNNLayer, self).__init__()
        self.W_msg = nn.Linear(ndim + edim, out_dim)       
        self.W_apply = nn.Linear(ndim + out_dim, out_dim)  
   
    def message_func(self, edges):
        return {'m': F.relu(self.W_msg(torch.cat([edges.src['h'], edges.data['h']], dim=1)))}
   
    def forward(self, g_dgl, nfeats, efeats):
        with g_dgl.local_scope():
            g = g_dgl
            g.ndata['h'] = nfeats 
            g.edata['h'] = efeats  
            g.update_all(self.message_func, fn.sum(msg='m', out='h_neigh'))
            g.ndata['h'] = F.relu(self.W_apply(torch.cat([g.ndata['h'], g.ndata['h_neigh']], dim=1)))
            return g.ndata['h']
   
   
    class GNN(nn.Module):
        def __init__(self, ndim, edim, out_dim, n_classes, dropout):
            super(GNN, self).__init__()
            self.layers = nn.ModuleList()
            self.layers.append(GNNLayer(ndim, edim, 50))
            self.layers.append(GNNLayer(50, edim, 25))
            self.layers.append(GNNLayer(25, edim, out_dim))
            self.fc = nn.Linear(out_dim, n_classes)   # added
            self.dropout = nn.Dropout(p=dropout)
   
        def forward(self, g, nfeats, efeats):
            for i, layer in enumerate(self.layers):
                if i != 0:
                    nfeats = self.dropout(nfeats)
                nfeats = layer(g, nfeats, efeats)
            g.ndata['h'] = nfeats                  # added
            h = dgl.mean_nodes(g, 'h')             # added
            return self.fc(h)                      # added
   

2. After the forward function, the returned g.ndata[‘h’] has incorparated both edge feature and node feature. Is that right (it’s a little confusing when the result is assigned to ndata[‘h’])?

I am just asking whether I’m on the right track. Thanks for your patience.

[mufeili]

if I want to adjust the codes to do graph classification (n_classes=3), is it okay to code like this? (Because I want to use crossentropy loss with softmax so I just add a fully-connected layer on top of the original model)

The code looks good to me.

After the forward function, the returned g.ndata[‘h’] has incorparated both edge feature and node feature. Is that right (it’s a little confusing when the result is assigned to ndata[‘h’])?

The output is self.fc(h), not g.ndata['h'], right?

[Silas]

Thank you! I got it.