dmlc--dgl
Dylan
b84de903a2
Correction like mentioned in #4969 I noticed that there is a normalisation step on line 97 while the normalised values are not used downstream. Even if this was meant to show the normalisation step, it would not be calculating the normalisation step described in the CGN paper. The paper considers both in and out degrees while the normalisation in the code only describes normalisation using the in degrees. In the end, the normalised values are assigned to g.ndata["norm"] but these values are not used afterwards. Having a normalisation step here is also unnecessary since the GraphConv layer that is used already takes care of the normalisation. https://docs.dgl.ai/en/0.9.x/_modules/dgl/nn/pytorch/conv/graphconv.html#GraphConv It confused me for a second thinking that I had to do the normalisation myself but this is already handled by the GraphConf.
110 行
3.2 KiB
Python
110 行
3.2 KiB
Python
import argparse
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
import dgl
|
|
import dgl.nn as dglnn
|
|
from dgl import AddSelfLoop
|
|
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
|
|
|
|
|
|
class GCN(nn.Module):
|
|
def __init__(self, in_size, hid_size, out_size):
|
|
super().__init__()
|
|
self.layers = nn.ModuleList()
|
|
# two-layer GCN
|
|
self.layers.append(
|
|
dglnn.GraphConv(in_size, hid_size, activation=F.relu)
|
|
)
|
|
self.layers.append(dglnn.GraphConv(hid_size, out_size))
|
|
self.dropout = nn.Dropout(0.5)
|
|
|
|
def forward(self, g, features):
|
|
h = features
|
|
for i, layer in enumerate(self.layers):
|
|
if i != 0:
|
|
h = self.dropout(h)
|
|
h = layer(g, h)
|
|
return h
|
|
|
|
|
|
def evaluate(g, features, labels, mask, model):
|
|
model.eval()
|
|
with torch.no_grad():
|
|
logits = model(g, features)
|
|
logits = logits[mask]
|
|
labels = labels[mask]
|
|
_, indices = torch.max(logits, dim=1)
|
|
correct = torch.sum(indices == labels)
|
|
return correct.item() * 1.0 / len(labels)
|
|
|
|
|
|
def train(g, features, labels, masks, model):
|
|
# define train/val samples, loss function and optimizer
|
|
train_mask = masks[0]
|
|
val_mask = masks[1]
|
|
loss_fcn = nn.CrossEntropyLoss()
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
|
|
|
|
# training loop
|
|
for epoch in range(200):
|
|
model.train()
|
|
logits = model(g, features)
|
|
loss = loss_fcn(logits[train_mask], labels[train_mask])
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
acc = evaluate(g, features, labels, val_mask, model)
|
|
print(
|
|
"Epoch {:05d} | Loss {:.4f} | Accuracy {:.4f} ".format(
|
|
epoch, loss.item(), acc
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--dataset",
|
|
type=str,
|
|
default="cora",
|
|
help="Dataset name ('cora', 'citeseer', 'pubmed').",
|
|
)
|
|
args = parser.parse_args()
|
|
print(f"Training with DGL built-in GraphConv module.")
|
|
|
|
# load and preprocess dataset
|
|
transform = (
|
|
AddSelfLoop()
|
|
) # by default, it will first remove self-loops to prevent duplication
|
|
if args.dataset == "cora":
|
|
data = CoraGraphDataset(transform=transform)
|
|
elif args.dataset == "citeseer":
|
|
data = CiteseerGraphDataset(transform=transform)
|
|
elif args.dataset == "pubmed":
|
|
data = PubmedGraphDataset(transform=transform)
|
|
else:
|
|
raise ValueError("Unknown dataset: {}".format(args.dataset))
|
|
g = data[0]
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
g = g.int().to(device)
|
|
features = g.ndata["feat"]
|
|
labels = g.ndata["label"]
|
|
masks = g.ndata["train_mask"], g.ndata["val_mask"], g.ndata["test_mask"]
|
|
|
|
# create GCN model
|
|
in_size = features.shape[1]
|
|
out_size = data.num_classes
|
|
model = GCN(in_size, 16, out_size).to(device)
|
|
|
|
# model training
|
|
print("Training...")
|
|
train(g, features, labels, masks, model)
|
|
|
|
# test the model
|
|
print("Testing...")
|
|
acc = evaluate(g, features, labels, masks[2], model)
|
|
print("Test accuracy {:.4f}".format(acc))
|