dmlc--dgl
maqy1995
280 行
11 KiB
Python
280 行
11 KiB
Python
import dgl
|
|
from functools import partial
|
|
import numpy as np
|
|
import torch as th
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torch.optim as optim
|
|
import torch.multiprocessing as mp
|
|
from torch.utils.data import DataLoader
|
|
import dgl.function as fn
|
|
import dgl.nn.pytorch as dglnn
|
|
import time
|
|
import argparse
|
|
from _thread import start_new_thread
|
|
from functools import wraps
|
|
from dgl.data import RedditDataset
|
|
import tqdm
|
|
import traceback
|
|
from ogb.nodeproppred import DglNodePropPredDataset
|
|
|
|
from sampler import ClusterIter, subgraph_collate_fn
|
|
|
|
#### Neighbor sampler
|
|
|
|
|
|
class GAT(nn.Module):
|
|
def __init__(self,
|
|
in_feats,
|
|
num_heads,
|
|
n_hidden,
|
|
n_classes,
|
|
n_layers,
|
|
activation,
|
|
dropout=0.):
|
|
super().__init__()
|
|
self.n_layers = n_layers
|
|
self.n_hidden = n_hidden
|
|
self.n_classes = n_classes
|
|
self.layers = nn.ModuleList()
|
|
self.num_heads = num_heads
|
|
self.layers.append(dglnn.GATConv(in_feats,
|
|
n_hidden,
|
|
num_heads=num_heads,
|
|
feat_drop=dropout,
|
|
attn_drop=dropout,
|
|
activation=activation,
|
|
negative_slope=0.2))
|
|
for i in range(1, n_layers - 1):
|
|
self.layers.append(dglnn.GATConv(n_hidden * num_heads,
|
|
n_hidden,
|
|
num_heads=num_heads,
|
|
feat_drop=dropout,
|
|
attn_drop=dropout,
|
|
activation=activation,
|
|
negative_slope=0.2))
|
|
self.layers.append(dglnn.GATConv(n_hidden * num_heads,
|
|
n_classes,
|
|
num_heads=num_heads,
|
|
feat_drop=dropout,
|
|
attn_drop=dropout,
|
|
activation=None,
|
|
negative_slope=0.2))
|
|
|
|
def forward(self, g, x):
|
|
h = x
|
|
for l, conv in enumerate(self.layers):
|
|
h = conv(g, h)
|
|
if l < len(self.layers) - 1:
|
|
h = h.flatten(1)
|
|
h = h.mean(1)
|
|
return h.log_softmax(dim=-1)
|
|
|
|
def inference(self, g, x, batch_size, device):
|
|
"""
|
|
Inference with the GAT model on full neighbors (i.e. without neighbor sampling).
|
|
g : the entire graph.
|
|
x : the input of entire node set.
|
|
The inference code is written in a fashion that it could handle any number of nodes and
|
|
layers.
|
|
"""
|
|
num_heads = self.num_heads
|
|
nodes = th.arange(g.number_of_nodes())
|
|
for l, layer in enumerate(self.layers):
|
|
if l < self.n_layers - 1:
|
|
y = th.zeros(g.number_of_nodes(), self.n_hidden * num_heads if l != len(self.layers) - 1 else self.n_classes)
|
|
else:
|
|
y = th.zeros(g.number_of_nodes(), self.n_hidden if l != len(self.layers) - 1 else self.n_classes)
|
|
sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
|
|
dataloader = dgl.dataloading.NodeDataLoader(
|
|
g,
|
|
th.arange(g.number_of_nodes()),
|
|
sampler,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
drop_last=False,
|
|
num_workers=args.num_workers)
|
|
|
|
for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
|
|
block = blocks[0].int().to(device)
|
|
h = x[input_nodes].to(device)
|
|
h_dst = h[:block.number_of_dst_nodes()].to(device)
|
|
if l < self.n_layers - 1:
|
|
h = layer(block, h).flatten(1)
|
|
else:
|
|
h = layer(block, h)
|
|
h = h.mean(1)
|
|
h = h.log_softmax(dim=-1)
|
|
|
|
y[output_nodes] = h.cpu()
|
|
x = y
|
|
return y
|
|
|
|
def compute_acc(pred, labels):
|
|
"""
|
|
Compute the accuracy of prediction given the labels.
|
|
"""
|
|
return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred)
|
|
|
|
def evaluate(model, g, labels, val_nid, test_nid, batch_size, device):
|
|
"""
|
|
Evaluate the model on the validation set specified by ``val_mask``.
|
|
g : The entire graph.
|
|
inputs : The features of all the nodes.
|
|
labels : The labels of all the nodes.
|
|
val_mask : A 0-1 mask indicating which nodes do we actually compute the accuracy for.
|
|
batch_size : Number of nodes to compute at the same time.
|
|
device : The GPU device to evaluate on.
|
|
"""
|
|
model.eval()
|
|
with th.no_grad():
|
|
inputs = g.ndata['feat']
|
|
pred = model.inference(g, inputs, batch_size, device)
|
|
model.train()
|
|
return compute_acc(pred[val_nid], labels[val_nid]), compute_acc(pred[test_nid], labels[test_nid]), pred
|
|
|
|
def model_param_summary(model):
|
|
""" Count the model parameters """
|
|
cnt = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
print("Total Params {}".format(cnt))
|
|
|
|
#### Entry point
|
|
def run(args, device, data):
|
|
# Unpack data
|
|
train_nid, val_nid, test_nid, in_feats, labels, n_classes, g, cluster_iterator = data
|
|
|
|
|
|
# Define model and optimizer
|
|
model = GAT(in_feats, args.num_heads, args.num_hidden, n_classes, args.num_layers, F.relu, args.dropout)
|
|
model_param_summary(model)
|
|
model = model.to(device)
|
|
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
|
|
|
|
# Training loop
|
|
avg = 0
|
|
best_eval_acc = 0
|
|
best_test_acc = 0
|
|
for epoch in range(args.num_epochs):
|
|
iter_load = 0
|
|
iter_far = 0
|
|
iter_back = 0
|
|
tic = time.time()
|
|
|
|
# Loop over the dataloader to sample the computation dependency graph as a list of
|
|
# blocks.
|
|
tic_start = time.time()
|
|
for step, cluster in enumerate(cluster_iterator):
|
|
cluster = cluster.int().to(device)
|
|
mask = cluster.ndata['train_mask']
|
|
if mask.sum() == 0:
|
|
continue
|
|
feat = cluster.ndata['feat']
|
|
batch_labels = cluster.ndata['labels']
|
|
tic_step = time.time()
|
|
|
|
# Compute loss and prediction
|
|
batch_pred = model(cluster, feat)
|
|
batch_pred = batch_pred[mask]
|
|
batch_labels = batch_labels[mask]
|
|
loss = nn.functional.nll_loss(batch_pred, batch_labels)
|
|
optimizer.zero_grad()
|
|
tic_far = time.time()
|
|
loss.backward()
|
|
optimizer.step()
|
|
tic_back = time.time()
|
|
iter_load += (tic_step - tic_start)
|
|
iter_far += (tic_far - tic_step)
|
|
iter_back += (tic_back - tic_far)
|
|
|
|
if step % args.log_every == 0:
|
|
acc = compute_acc(batch_pred, batch_labels)
|
|
gpu_mem_alloc = th.cuda.max_memory_allocated() / 1000000 if th.cuda.is_available() else 0
|
|
print('Epoch {:05d} | Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | GPU {:.1f} MB'.format(
|
|
epoch, step, loss.item(), acc.item(), gpu_mem_alloc))
|
|
tic_start = time.time()
|
|
|
|
toc = time.time()
|
|
print('Epoch Time(s): {:.4f} Load {:.4f} Forward {:.4f} Backward {:.4f}'.format(toc - tic, iter_load, iter_far, iter_back))
|
|
if epoch >= 5:
|
|
avg += toc - tic
|
|
|
|
if epoch % args.eval_every == 0 and epoch != 0:
|
|
eval_acc, test_acc, pred = evaluate(model, g, labels, val_nid, test_nid, args.val_batch_size, device)
|
|
model = model.to(device)
|
|
if args.save_pred:
|
|
np.savetxt(args.save_pred + '%02d' % epoch, pred.argmax(1).cpu().numpy(), '%d')
|
|
print('Eval Acc {:.4f}'.format(eval_acc))
|
|
if eval_acc > best_eval_acc:
|
|
best_eval_acc = eval_acc
|
|
best_test_acc = test_acc
|
|
print('Best Eval Acc {:.4f} Test Acc {:.4f}'.format(best_eval_acc, best_test_acc))
|
|
print('Avg epoch time: {}'.format(avg / (epoch - 4)))
|
|
return best_test_acc
|
|
|
|
if __name__ == '__main__':
|
|
argparser = argparse.ArgumentParser("multi-gpu training")
|
|
argparser.add_argument('--gpu', type=int, default=0,
|
|
help="GPU device ID. Use -1 for CPU training")
|
|
argparser.add_argument('--num-epochs', type=int, default=20)
|
|
argparser.add_argument('--num-hidden', type=int, default=128)
|
|
argparser.add_argument('--num-layers', type=int, default=3)
|
|
argparser.add_argument('--num-heads', type=int, default=8)
|
|
argparser.add_argument('--batch-size', type=int, default=32)
|
|
argparser.add_argument('--val-batch-size', type=int, default=2000)
|
|
argparser.add_argument('--log-every', type=int, default=20)
|
|
argparser.add_argument('--eval-every', type=int, default=1)
|
|
argparser.add_argument('--lr', type=float, default=0.001)
|
|
argparser.add_argument('--dropout', type=float, default=0.5)
|
|
argparser.add_argument('--save-pred', type=str, default='')
|
|
argparser.add_argument('--wd', type=float, default=0)
|
|
argparser.add_argument('--num_partitions', type=int, default=15000)
|
|
argparser.add_argument('--num-workers', type=int, default=0)
|
|
args = argparser.parse_args()
|
|
|
|
if args.gpu >= 0:
|
|
device = th.device('cuda:%d' % args.gpu)
|
|
else:
|
|
device = th.device('cpu')
|
|
|
|
# load reddit data
|
|
data = DglNodePropPredDataset(name='ogbn-products')
|
|
splitted_idx = data.get_idx_split()
|
|
train_idx, val_idx, test_idx = splitted_idx['train'], splitted_idx['valid'], splitted_idx['test']
|
|
graph, labels = data[0]
|
|
labels = labels[:, 0]
|
|
print('Total edges before adding self-loop {}'.format(graph.number_of_edges()))
|
|
graph = dgl.remove_self_loop(graph)
|
|
graph = dgl.add_self_loop(graph)
|
|
print('Total edges after adding self-loop {}'.format(graph.number_of_edges()))
|
|
num_nodes = train_idx.shape[0] + val_idx.shape[0] + test_idx.shape[0]
|
|
assert num_nodes == graph.number_of_nodes()
|
|
graph.ndata['labels'] = labels
|
|
mask = th.zeros(num_nodes, dtype=th.bool)
|
|
mask[train_idx] = True
|
|
graph.ndata['train_mask'] = mask
|
|
mask = th.zeros(num_nodes, dtype=th.bool)
|
|
mask[val_idx] = True
|
|
graph.ndata['valid_mask'] = mask
|
|
mask = th.zeros(num_nodes, dtype=th.bool)
|
|
mask[test_idx] = True
|
|
graph.ndata['test_mask'] = mask
|
|
|
|
graph.in_degrees(0)
|
|
graph.out_degrees(0)
|
|
graph.find_edges(0)
|
|
|
|
cluster_iter_data = ClusterIter(
|
|
'ogbn-products', graph, args.num_partitions, args.batch_size)
|
|
cluster_iterator = DataLoader(cluster_iter_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=4, collate_fn=partial(subgraph_collate_fn, graph))
|
|
|
|
in_feats = graph.ndata['feat'].shape[1]
|
|
n_classes = (labels.max() + 1).item()
|
|
# Pack data
|
|
data = train_idx, val_idx, test_idx, in_feats, labels, n_classes, graph, cluster_iterator
|
|
|
|
# Run 10 times
|
|
test_accs = []
|
|
for i in range(10):
|
|
test_accs.append(run(args, device, data))
|
|
print('Average test accuracy:', np.mean(test_accs), '±', np.std(test_accs))
|