Skip to content

Instantly share code, notes, and snippets.

Created Sep 26, 2021
What would you like to do?
import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F
# Generate a synthetic dataset with 10000 graphs, ranging from 10 to 500 nodes.
dataset ='PROTEINS', self_loop=True)
print('Node feature dimensionality:', dataset.dim_nfeats)
print('Number of graph categories:', dataset.gclasses)
from dgl.dataloading import GraphDataLoader
from import SubsetRandomSampler
num_examples = len(dataset)
num_train = int(num_examples * 0.8)
train_sampler = SubsetRandomSampler(torch.arange(num_train))
test_sampler = SubsetRandomSampler(torch.arange(num_train, num_examples))
train_dataloader = GraphDataLoader(
dataset, sampler=train_sampler, batch_size=5, drop_last=False)
test_dataloader = GraphDataLoader(
dataset, sampler=test_sampler, batch_size=5, drop_last=False)
it = iter(train_dataloader)
batch = next(it)
batched_graph, labels = batch
print('Number of nodes for each graph element in the batch:', batched_graph.batch_num_nodes())
print('Number of edges for each graph element in the batch:', batched_graph.batch_num_edges())
# Recover the original graph elements from the minibatch
graphs = dgl.unbatch(batched_graph)
print('The original graphs in the minibatch:')
from dgl.nn import GraphConv, GATConv
class GCN(nn.Module):
def __init__(self, in_feats, h_feats, num_classes):
super(GCN, self).__init__()
self.conv1 = GraphConv(in_feats, h_feats)
self.conv2 = GraphConv(h_feats, num_classes)
def forward(self, g, in_feat):
h = self.conv1(g, in_feat)
h = F.relu(h)
h = self.conv2(g, h)
g.ndata['h'] = h
x = dgl.mean_nodes(g, 'h')
print("--------- x.shape: {}, g: {}".format(x.shape, g))
return x
class GAT(nn.Module):
def __init__(self, in_feats, h_feats, num_classes, num_heads=3):
super(GAT, self).__init__()
self.conv1 = GATConv(in_feats, h_feats, num_heads=num_heads)
self.conv2 = GATConv(h_feats, h_feats, num_heads)
self.conv3 = torch.nn.Linear(in_features=h_feats, out_features=num_classes)
def forward(self, g, in_feat):
h = self.conv1(g, in_feat)
h = torch.mean(h, dim=1)
h = self.conv2(g, h)
h = torch.mean(h, dim=1)
h = self.conv3(h)
g.ndata['h'] = h
x = dgl.mean_nodes(g, 'h')
print("--------- h.shape: {}, x.shape: {}, g.batch_size: {}".format(h.shape, x.shape, g.batch_size))
return x
# Create the model with given dimensions
#model = GCN(dataset.dim_nfeats, 16, dataset.gclasses)
model = GAT(dataset.dim_nfeats, 16, dataset.gclasses)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(1):
for batched_graph, labels in train_dataloader:
pred = model(batched_graph, batched_graph.ndata['attr'].float())
loss = F.cross_entropy(pred, labels)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment