DGL GCN Model

dgl-gcn-model.py

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

class GCN(nn.Module):
    def __init__(self, in_feats, h_feats):
        super(GCN, self).__init__()
        self.conv1 = GraphConv(in_feats, h_feats)
        self.conv2 = GraphConv(h_feats, h_feats)

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = F.relu(self.conv1(g, in_feat))
        h = F.dropout(h, training=self.training)
        h = self.conv2(g, h)
        return F.log_softmax(h, dim=1)

node_features = data.ndata['feat']
node_labels = data.ndata['label']
train_mask = data.ndata['train_mask']
test_mask = data.ndata['test_mask']
n_features = node_features.shape[1]

# instantiate the model
model = GCN(n_features, 16)