def mol_to_nearest_neighbor_graph(mol,
coordinates,
neighbor_cutoff,
max_num_neighbors=None,
p_distance=2,
add_self_loop=False,
node_featurizer=None,
edge_featurizer=None,
canonical_atom_order=True,
keep_dists=False,
dist_field='dist'):
"""Convert an RDKit molecule into a nearest neighbor graph and featurize for it.
Different from bigraph and complete graph, the nearest neighbor graph
may not be symmetric since i is the closest neighbor of j does not
necessarily suggest the other way.
Parameters
----------
mol : rdkit.Chem.rdchem.Mol
RDKit molecule holder
coordinates : numpy.ndarray of shape (N, D)
The coordinates of atoms in the molecule. N for the number of atoms
and D for the dimensions of the coordinates.
neighbor_cutoff : float
If the distance between a pair of nodes is larger than neighbor_cutoff,
they will not be considered as neighboring nodes.
max_num_neighbors : int or None.
If not None, then this specifies the maximum number of neighbors
allowed for each atom. Default to None.
p_distance : int
We compute the distance between neighbors using Minkowski (:math:`l_p`)
distance. When ``p_distance = 1``, Minkowski distance is equivalent to
Manhattan distance. When ``p_distance = 2``, Minkowski distance is
equivalent to the standard Euclidean distance. Default to 2.
add_self_loop : bool
Whether to add self loops in DGLGraphs. Default to False.
node_featurizer : callable, rdkit.Chem.rdchem.Mol -> dict
Featurization for nodes like atoms in a molecule, which can be used to update
ndata for a DGLGraph. Default to None.
edge_featurizer : callable, rdkit.Chem.rdchem.Mol -> dict
Featurization for edges like bonds in a molecule, which can be used to update
edata for a DGLGraph. Default to None.
canonical_atom_order : bool
Whether to use a canonical order of atoms returned by RDKit. Setting it
to true might change the order of atoms in the graph constructed. Default
to True.
keep_dists : bool
Whether to store the distance between neighboring atoms in ``edata`` of the
constructed DGLGraphs. Default to False.
dist_field : str
Field for storing distance between neighboring atoms in ``edata``. This comes
into effect only when ``keep_dists=True``. Default to ``'dist'``.
"""
if canonical_atom_order:
new_order = rdmolfiles.CanonicalRankAtoms(mol)
mol = rdmolops.RenumberAtoms(mol, new_order)
srcs, dsts, dists = k_nearest_neighbors(
coordinates=coordinates,
neighbor_cutoff=neighbor_cutoff,
max_num_neighbors=max_num_neighbors,
p_distance=p_distance,
self_loops=add_self_loop)
g = DGLGraph()
# Add nodes first since some nodes may be completely isolated
num_atoms = mol.GetNumAtoms()
g.add_nodes(num_atoms)
# Add edges
g.add_edges(srcs, dsts)
if node_featurizer is not None:
g.ndata.update(node_featurizer(mol))
if edge_featurizer is not None:
g.edata.update(edge_featurizer(mol))
if keep_dists:
assert dist_field not in g.edata, \
'Expect {} to be reserved for distance between neighboring atoms.'
g.edata[dist_field] = torch.tensor(dists).float().reshape(-1, 1)
return g
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl import DGLGraph
from dgl.nn.pytorch import GraphConv
from preprocess import adj
import numpy as np
import time as time
if __name__ == '__main__':
g = DGLGraph()
g.from_scipy_sparse_matrix(adj)
n = g.number_of_nodes()
gc = GraphConv(1433, 16, norm=False, bias=False)
gc.cuda()
gc.eval()
h0 = torch.randn(n, 1433).cuda()
times = []
for i in range(10):
t_beg = time.time()
h1 = gc(h0, g)
t_end = time.time()
print('time = {} s'.format(t_end - t_beg))
times.append(t_end - t_beg)
print('Average time = {} s'.format(np.mean(times)))
def main(args):
# load and preprocess dataset
# data = load_dgl_data(args)
dataset = args.dataset
# prefix = '/mnt/yushi/'
# prefix = 'graphzoom'
dataset_dir = f'{args.prefix}/dataset/{dataset}'
# data = load_data(dataset_dir, args.dataset)
load_data_time = time.time()
# if dataset in ['Amazon2M', 'reddit']:
if dataset in ['Amazon2M']:
g, _ = load_graphs(
f'{args.prefix}/dataset/Amazon2M/Amazon2M_dglgraph.bin')
g = g[0]
data = g.ndata
features = torch.FloatTensor(data['feat'])
onehot_labels = F.one_hot(data['label']).numpy()
train_mask = data['train_mask'].bool()
val_mask = data['val_mask'].bool()
test_mask = val_mask
data = EasyDict({
'graph': g,
'labels': data['label'],
'onehot_labels': onehot_labels,
'features': data['feat'],
'train_mask': train_mask,
'val_mask': val_mask,
'test_mask': test_mask,
'num_labels': onehot_labels.shape[1],
'coarse': False
})
else:
original_adj, labels, train_ids, test_ids, train_labels, test_labels, feats = load_data(
dataset_dir, args.dataset)
data = load_dgl_data(args)
labels = torch.LongTensor(labels)
train_mask = _sample_mask(train_ids, labels.shape[0])
onehot_labels = F.one_hot(labels).numpy()
if dataset == 'reddit':
g = data.graph
else:
val_ids = test_ids[1000:1500]
test_ids = test_ids[:1000]
test_mask = _sample_mask(test_ids, labels.shape[0])
val_mask = _sample_mask(val_ids, labels.shape[0])
data = EasyDict({
'graph': data.graph,
'labels': labels,
'onehot_labels': onehot_labels,
'features': feats,
# 'features': data.features,
'train_mask': train_mask,
'val_mask': val_mask,
'test_mask': test_mask,
'num_labels': onehot_labels.shape[1],
'coarse': False
})
# g = DGLGraph(data.graph)
print(f'load data finished: {time.time() - load_data_time}')
if args.coarse:
# * load projection matrix
levels = args.level
reduce_results = f"graphzoom/reduction_results/{dataset}/fusion/"
projections, coarse_adj = construct_proj_laplacian(
original_adj, levels, reduce_results)
# *calculate coarse feature, labels
label_mask = np.expand_dims(data.train_mask, 1)
onehot_labels = onehot_labels * label_mask
for i in range(levels):
data.features = projections[i] @ data.features
onehot_labels = projections[i] @ onehot_labels
# coarse_labels = projections[0] @ onehot_labels
# ! add train_mask
rows_sum = onehot_labels.sum(axis=1)[:, np.newaxis]
norm_coarse_labels = onehot_labels / rows_sum
norm_label_entropy = Categorical(
torch.Tensor(norm_coarse_labels)).entropy()
label_entropy_mask = torch.BoolTensor(norm_label_entropy < 0.01)
coarse_train_mask = torch.BoolTensor(onehot_labels.sum(axis=1))
# coarse_train_mask = label_entropy_mask
# ! entropy threshold
# coarse_graph = nx.Graph(coarse_adj[1])
print('creating coarse DGLGraph')
start = time.process_time()
g = DGLGraph()
g.from_scipy_sparse_matrix(coarse_adj[1])
print(f'creating finished in {time.process_time() - start}')
# list(map(np.shape, [coarse_embed, coarse_labels]))
# * new train/test masks
coarsen_ratio = projections[0].shape[1] / projections[0].shape[0]
# coarse_train_mask = _sample_mask(
# range(int(coarsen_ratio*len(data.train_mask.int().sum().item()))),
# norm_coarse_labels.shape[0])
# coarse_train_mask = _sample_mask(
# range(norm_coarse_labels.shape[0]),
# range(60),
# norm_coarse_labels.shape[0])
# coarse_test_mask = _sample_mask(
# range(100, 700), norm_coarse_labels.shape[0])
# coarse_val_mask = _sample_mask(
# range(700, 1000), norm_coarse_labels.shape[0])
# *replace data
data = EasyDict({
'graph': g,
'labels': onehot_labels,
# 'onehot_labels': onehot_labels,
'features': data.features,
'train_mask': coarse_train_mask,
# 'val_mask': coarse_val_mask,
# 'test_mask': coarse_test_mask,
'num_classes': norm_coarse_labels.shape[1],
'num_labels': onehot_labels.shape[1],
'coarse': True
})
if args.coarse:
labels = torch.FloatTensor(data.labels)
loss_fcn = torch.nn.KLDivLoss(reduction='batchmean')
print('training coarse')
else:
labels = torch.LongTensor(data.labels)
loss_fcn = torch.nn.CrossEntropyLoss()
features = torch.FloatTensor(data.features)
if hasattr(torch, 'BoolTensor'):
train_mask = torch.BoolTensor(data.train_mask)
val_mask = torch.BoolTensor(data.val_mask)
test_mask = torch.BoolTensor(data.test_mask)
else:
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
# graph preprocess and calculate normalization factor
# add self loop
if args.self_loop or args.arch == 'gat':
g = add_self_loop(data.graph)
print('add self_loop')
n_edges = g.number_of_edges()
print("""----Data statistics------'
# Edges %d
# Classes %d
# Train samples %d
# Val samples %d
# Test samples %d""" %
(n_edges, n_classes, train_mask.int().sum().item(),
val_mask.int().sum().item(), test_mask.int().sum().item()))
# normalization
degs = g.in_degrees().float()
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
if cuda:
norm = norm.cuda()
g.ndata['norm'] = norm.unsqueeze(1)
# * create GCN model
heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
model = create_model(
args.arch,
g,
num_layers=args.num_layers,
in_dim=in_feats,
num_hidden=args.num_hidden,
num_classes=n_classes,
heads=heads,
# activation=F.elu,
feat_drop=args.in_drop,
attn_drop=args.attn_drop,
negative_slope=args.negative_slope,
residual=args.residual,
log_softmax=args.coarse)
if cuda:
model.cuda()
print(model)
# loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
acc = 0
start = time.time()
for epoch in range(args.n_epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits, h = model(features)
loss = loss_fcn(logits[train_mask], labels[train_mask]) # ?
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
if not args.coarse:
acc = evaluate(model, features, labels, val_mask)
print(
"Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
"ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item(),
acc, n_edges / np.mean(dur) / 1000))
print(f'training time: {time.time() - start}')
if not args.coarse:
acc = evaluate(model, features, labels, test_mask)
print(h.shape)
np.save(f'embeddings/{(args.arch).upper()}_{dataset}_emb_level_1_mask',
h.detach().cpu().numpy())
torch.save(
model.state_dict(),
f'embeddings/{(args.arch).upper()}_{dataset}_emb_level_1_params.pth.tar',
)
print("Test accuracy {:.2%}".format(acc))
def train_gcn(dataset,
test_ratio=0.5,
val_ratio=0.2,
seed=1,
n_hidden=64,
n_epochs=200,
lr=1e-2,
weight_decay=5e-4,
dropout=0.5,
verbose=True,
cuda=False):
data = dataset.get_data()
features = torch.FloatTensor(data['features'])
labels = torch.LongTensor(data['labels'])
n = len(data['ids'])
train_mask, val_mask, test_mask = get_masks(n,
data['main_ids'],
data['main_labels'],
test_ratio=test_ratio,
val_ratio=val_ratio,
seed=seed)
train_mask = torch.BoolTensor(train_mask)
val_mask = torch.BoolTensor(val_mask)
test_mask = torch.BoolTensor(test_mask)
if cuda:
torch.cuda.set_device("cuda:0")
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
g = DGLGraph(data['graph'])
g = dgl.transform.add_self_loop(g)
n_edges = g.number_of_edges()
degs = g.in_degrees().float()
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
if cuda:
norm = norm.cuda()
g.ndata['norm'] = norm.unsqueeze(1)
in_feats = features.shape[1]
# + 1 for unknown class
n_classes = data['n_classes'] + 1
model = GCN(g,
in_feats=in_feats,
n_hidden=n_hidden,
n_classes=n_classes,
activation=F.relu,
dropout=dropout)
if cuda:
model.cuda()
loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode='min',
factor=0.9,
patience=20,
min_lr=1e-10)
best_f1 = -100
# initialize graph
dur = []
for epoch in range(n_epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(features)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
f1 = evaluate(model, features, labels, val_mask)
scheduler.step(1 - f1)
if f1 > best_f1:
best_f1 = f1
torch.save(model.state_dict(), 'best_model.pt')
if verbose:
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | F1 {:.4f} | "
"ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur),
loss.item(), f1,
n_edges / np.mean(dur) / 1000))
model.load_state_dict(torch.load('best_model.pt'))
f1 = evaluate(model, features, labels, test_mask)
if verbose:
print()
print("Test F1 {:.2}".format(f1))
return f1
#3M_276M = citation_graph.load_RMAT('3M_276M',100,10)
_21M = citation_graph.load_RMAT('21', 1024, 10)
#_22M = citation_graph.load_RMAT('22',100,10)
#_23M = citation_graph.load_RMAT('23',100,10)
#24M = citation_graph.load_RMAT('24',100,10)
#25M = citation_graph.load_RMAT('25',100,10)
#26M = citation_graph.load_RMAT('26',100,10)
#F_64 = citation_graph.load_RMAT('feature',64,16)
#AIFB = load_data('aifb', bfs_level=3)
#MUTAG = load_data('mutag', bfs_level=3)
#BGS = load_data('bgs', bfs_level=3)
#AM = load_data('am', bfs_level=3)
#One training run before we start tracking duration to warm up GPU.
g = DGLGraph(Cora.graph)
norm = torch.pow(g.in_degrees().float(), -0.5)
norm[torch.isinf(norm)] = 0
g.ndata['norm'] = norm.unsqueeze(1).to(device)
model = GCN(g, Cora.features.shape[1], Cora.num_labels).to(device)
inference(model, Cora, epochs=10, device=device)
print('---------start------------')
DatasetList = [_21M]
#ModelList = [GCN, GCN_GRU, GCN_GATED, GraphSAGE, SGConv, TAGCN]
ModelList = [GCN]
result = []
for d in DatasetList:
g = DGLGraph(d.graph)
norm = torch.pow(g.in_degrees().float(), -0.5)
norm[torch.isinf(norm)] = 0
g.ndata['norm'] = norm.unsqueeze(1).to(device)
def main(args):
# load and preprocess dataset
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" % (n_edges, n_classes, train_mask.sum().item(),
val_mask.sum().item(), test_mask.sum().item()))
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
# graph preprocess and calculate normalization factor
g = DGLGraph(data.graph)
n_edges = g.number_of_edges()
# add self loop
g.add_edges(g.nodes(), g.nodes())
# create SGC model
model = SGConv(in_feats, n_classes, k=2, cached=True, bias=args.bias)
if cuda: model.cuda()
loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
for epoch in range(args.n_epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(g, features) # only compute the train set
torch.cuda.synchronize()
if epoch >= 3:
dur.append(time.time() - t0)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
acc = evaluate(model, g, features, labels, val_mask)
print(
"Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
"ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item(),
acc, n_edges / np.mean(dur) / 1000))
print()
acc = evaluate(model, g, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc))
def main(args):
# load and preprocess dataset
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
if hasattr(torch, 'BoolTensor'):
train_mask = torch.BoolTensor(data.train_mask)
val_mask = torch.BoolTensor(data.val_mask)
test_mask = torch.BoolTensor(data.test_mask)
else:
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
num_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes, train_mask.int().sum().item(),
val_mask.int().sum().item(), test_mask.int().sum().item()))
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
g = data.graph
# add self loop
g.remove_edges_from(nx.selfloop_edges(g))
g = DGLGraph(g)
g.add_edges(g.nodes(), g.nodes())
if cuda:
g = g.to(args.gpu)
n_edges = g.number_of_edges()
# create model
heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
model = GAT(g, args.num_layers, num_feats, args.num_hidden, n_classes,
heads, F.elu, args.in_drop, args.attn_drop,
args.negative_slope, args.residual)
print(model)
if args.early_stop:
stopper = EarlyStopping(patience=100)
if cuda:
model.cuda()
loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
for epoch in range(args.epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(features)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
train_acc = accuracy(logits[train_mask], labels[train_mask])
if args.fastmode:
val_acc = accuracy(logits[val_mask], labels[val_mask])
else:
val_acc = evaluate(model, features, labels, val_mask)
if args.early_stop:
if stopper.step(val_acc, model):
break
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | TrainAcc {:.4f} |"
" ValAcc {:.4f} | ETputs(KTEPS) {:.2f}".format(
epoch, np.mean(dur), loss.item(), train_acc, val_acc,
n_edges / np.mean(dur) / 1000))
print()
if args.early_stop:
model.load_state_dict(torch.load('es_checkpoint.pt'))
acc = evaluate(model, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc))
def test_graph6():
"""Batched graph with node types and edge distances."""
g1 = DGLGraph([(0, 1), (0, 2), (1, 2)])
g2 = DGLGraph([(0, 1), (1, 2), (1, 3), (1, 4)])
bg = dgl.batch([g1, g2])
return bg, torch.LongTensor([0, 1, 0, 2, 0, 3, 4, 4]), torch.randn(7, 1)
def test_graph7():
"""Graph with categorical node and edge features."""
g1 = DGLGraph([(0, 1), (0, 2), (1, 2)])
return g1, torch.LongTensor([0, 1, 0]), torch.LongTensor([2, 3, 4]), \
torch.LongTensor([0, 0, 1]), torch.LongTensor([2, 3, 2])
def test_graph3():
"""Graph with node features and edge features."""
g = DGLGraph([(0, 1), (0, 2), (1, 2)])
return g, torch.arange(g.number_of_nodes()).float().reshape(-1, 1), \
torch.arange(2 * g.number_of_edges()).float().reshape(-1, 2)
def test_graph5():
"""Graph with node types and edge distances."""
g1 = DGLGraph([(0, 1), (0, 2), (1, 2)])
return g1, torch.LongTensor([0, 1, 0]), torch.randn(3, 1)
def test_graph2():
"""Batched graph with node features."""
g1 = DGLGraph([(0, 1), (0, 2), (1, 2)])
g2 = DGLGraph([(0, 1), (1, 2), (1, 3), (1, 4)])
bg = dgl.batch([g1, g2])
return bg, torch.arange(bg.number_of_nodes()).float().reshape(-1, 1)
def main(args):
# load and preprocess dataset
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
if hasattr(torch, "BoolTensor"):
train_mask = torch.BoolTensor(data.train_mask)
val_mask = torch.BoolTensor(data.val_mask)
test_mask = torch.BoolTensor(data.test_mask)
else:
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" % (
n_edges,
n_classes,
train_mask.int().sum().item(),
val_mask.int().sum().item(),
test_mask.int().sum().item(),
))
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
# graph preprocess and calculate normalization factor
g = data.graph
# add self loop
if args.self_loop:
g.remove_edges_from(nx.selfloop_edges(g))
g.add_edges_from(zip(g.nodes(), g.nodes()))
g = DGLGraph(g)
n_edges = g.number_of_edges()
# normalization
degs = g.in_degrees().float()
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
if cuda:
norm = norm.cuda()
g.ndata["norm"] = norm.unsqueeze(1)
# create GCN model
model = GCN(
g,
in_feats,
args.n_hidden,
n_classes,
args.n_layers,
F.relu,
args.dropout,
)
if cuda:
model.cuda()
loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
for epoch in range(args.n_epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(features)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
accuracy, precision, recall, fscore, _ = evaluate(
model, features, labels, val_mask)
print("Epoch:", epoch)
print("Loss:", loss.item())
print("Accuracy:", accuracy)
print("Precision:", precision)
print("Recall:", recall)
print("F-Score:", fscore)
print()
print("=" * 80)
print()
accuracy, precision, recall, fscore, class_based_report = evaluate(
model, features, labels, test_mask)
print("=" * 80)
print(" " * 28 + "Final Statistics")
print("=" * 80)
print("Accuracy", accuracy)
print("Precision", precision)
print("Recall", recall)
print("F-Score", fscore)
print(class_based_report)
def main(args):
# load and preprocess dataset
data = load_data(args)
##
structure_features = np.load('../../pretrained/' + args.dataset +
'_structure_32d.npy')
attr_features = np.load('../../pretrained/' + args.dataset +
'_attr_32d.npy')
galpha = args.galpha
gbeta = args.gbeta
structure_features = preprocessing.scale(structure_features,
axis=1,
with_mean=True,
with_std=True,
copy=True)
#structure_features = preprocessing.scale(structure_features, axis=0, with_mean=True,with_std=True,copy=True)
structure_features = torch.FloatTensor(structure_features).cuda()
attr_features = preprocessing.scale(attr_features,
axis=1,
with_mean=True,
with_std=True,
copy=True)
#attr_features = preprocessing.scale(attr_features, axis=0, with_mean=True,with_std=True,copy=True)
attr_features = torch.FloatTensor(attr_features).cuda()
in_feats2 = structure_features.shape[1]
in_feats3 = attr_features.shape[1]
print(structure_features.shape, attr_features.shape)
##
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats1 = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" % (n_edges, n_classes, train_mask.sum().item(),
val_mask.sum().item(), test_mask.sum().item()))
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
# graph preprocess and calculate normalization factor
g = DGLGraph(data.graph)
n_edges = g.number_of_edges()
# add self loop
g.add_edges(g.nodes(), g.nodes())
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
# normalization
degs = g.in_degrees().float()
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
if cuda:
norm = norm.cuda()
g.ndata['norm'] = norm.unsqueeze(1)
# create APPNP model
#alpha2_set = [0,0.001,0.002,0.004,0.006,0.008,0.01,0.02,0.03,0.04,0.05]
#alpha3_set = [0,0.001,0.002,0.004,0.006,0.008,0.01,0.02,0.03,0.04,0.05]
result = []
for iter in range(10):
model = APPNP(g, in_feats1, in_feats2, in_feats3, args.hidden_sizes,
n_classes, F.relu, args.in_drop, args.edge_drop,
args.alpha, args.k, 1, galpha, gbeta)
if cuda:
model.cuda()
loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
best_val_acc = 0
best_test_acc = 0
for epoch in range(args.n_epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(features, structure_features, attr_features)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
val_acc = evaluate(model, features, structure_features,
attr_features, labels, val_mask)
if val_acc >= best_val_acc:
best_val_acc = val_acc
best_test_acc = evaluate(model, features, structure_features,
attr_features, labels, test_mask)
result.append(best_test_acc)
print('average result of 10 experiments:', np.average(result))
def test_send_multigraph():
g = DGLGraph(multigraph=True)
g.add_nodes(3)
g.add_edge(0, 1)
g.add_edge(0, 1)
g.add_edge(0, 1)
g.add_edge(2, 1)
def _message_a(edges):
return {'a': edges.data['a']}
def _message_b(edges):
return {'a': edges.data['a'] * 3}
def _reduce(nodes):
return {'a': F.max(nodes.mailbox['a'], 1)}
def answer(*args):
return F.max(F.stack(args, 0), 0)
# send by eid
old_repr = F.randn((4, 5))
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send([0, 2], message_func=_message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2]))
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send([0, 2, 3], message_func=_message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
# send on multigraph
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send(([0, 2], [1, 1]), _message_a)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], F.max(old_repr, 0))
# consecutive send and send_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send((2, 1), _message_a)
g.send([0, 1], message_func=_message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0] * 3, old_repr[1] * 3, old_repr[3]))
# consecutive send_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send(0, message_func=_message_a)
g.send(1, message_func=_message_b)
g.recv(1, _reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[1] * 3))
# send_and_recv_on
g.ndata['a'] = F.zeros((3, 5))
g.edata['a'] = old_repr
g.send_and_recv([0, 2, 3], message_func=_message_a, reduce_func=_reduce)
new_repr = g.ndata['a']
assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
assert F.allclose(new_repr[[0, 2]], F.zeros((2, 5)))
def main(args):
# load and preprocess dataset
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
if hasattr(torch, 'BoolTensor'):
train_mask = torch.BoolTensor(data.train_mask)
val_mask = torch.BoolTensor(data.val_mask)
test_mask = torch.BoolTensor(data.test_mask)
else:
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
# graph preprocess
g = data.graph
# add self loop
if args.self_loop:
g.remove_edges_from(nx.selfloop_edges(g))
g.add_edges_from(zip(g.nodes(), g.nodes()))
g = DGLGraph(g)
n_edges = g.number_of_edges()
if args.gpu >= 0:
g = g.to(args.gpu)
# create DGI model
dgi = DGI(g, in_feats, args.n_hidden, args.n_layers,
nn.PReLU(args.n_hidden), args.dropout)
if cuda:
dgi.cuda()
dgi_optimizer = torch.optim.Adam(dgi.parameters(),
lr=args.dgi_lr,
weight_decay=args.weight_decay)
# train deep graph infomax
cnt_wait = 0
best = 1e9
best_t = 0
dur = []
for epoch in range(args.n_dgi_epochs):
dgi.train()
if epoch >= 3:
t0 = time.time()
dgi_optimizer.zero_grad()
loss = dgi(features)
loss.backward()
dgi_optimizer.step()
if loss < best:
best = loss
best_t = epoch
cnt_wait = 0
torch.save(dgi.state_dict(), 'best_dgi.pkl')
else:
cnt_wait += 1
if cnt_wait == args.patience:
print('Early stopping!')
break
if epoch >= 3:
dur.append(time.time() - t0)
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | "
"ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item(),
n_edges / np.mean(dur) / 1000))
# create classifier model
classifier = Classifier(args.n_hidden, n_classes)
if cuda:
classifier.cuda()
classifier_optimizer = torch.optim.Adam(classifier.parameters(),
lr=args.classifier_lr,
weight_decay=args.weight_decay)
# train classifier
print('Loading {}th epoch'.format(best_t))
dgi.load_state_dict(torch.load('best_dgi.pkl'))
embeds = dgi.encoder(features, corrupt=False)
embeds = embeds.detach()
dur = []
for epoch in range(args.n_classifier_epochs):
classifier.train()
if epoch >= 3:
t0 = time.time()
classifier_optimizer.zero_grad()
preds = classifier(embeds)
loss = F.nll_loss(preds[train_mask], labels[train_mask])
loss.backward()
classifier_optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
acc = evaluate(classifier, embeds, labels, val_mask)
print(
"Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
"ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item(),
acc, n_edges / np.mean(dur) / 1000))
print()
acc = evaluate(classifier, embeds, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc))
def main(args):
# load graph data
data = load_data(args.dataset, bfs_level=args.bfs_level, relabel=args.relabel)
num_nodes = data.num_nodes
num_rels = data.num_rels
num_classes = data.num_classes
labels = data.labels
train_idx = data.train_idx
test_idx = data.test_idx
# split dataset into train, validate, test
if args.validation:
val_idx = train_idx[:len(train_idx) // 5]
train_idx = train_idx[len(train_idx) // 5:]
else:
val_idx = train_idx
# since the nodes are featureless, the input feature is then the node id.
feats = torch.arange(num_nodes)
# edge type and normalization factor
edge_type = torch.from_numpy(data.edge_type).long()
edge_norm = torch.from_numpy(data.edge_norm).unsqueeze(1).long()
labels = torch.from_numpy(labels).view(-1).long()
# check cuda
use_cuda = args.gpu >= 0 and torch.cuda.is_available()
if use_cuda:
torch.cuda.set_device(args.gpu)
feats = feats.cuda()
edge_type = edge_type.cuda()
edge_norm = edge_norm.cuda()
labels = labels.cuda()
# create graph
g = DGLGraph()
g.add_nodes(num_nodes)
g.add_edges(data.edge_src, data.edge_dst)
#tu_forward = sorted(list(zip(data.edge_src, data.edge_dst, data.edge_type)), key=lambda x : (x[1], x[2]))
#tu_backward = sorted(list(zip(data.edge_dst, data.edge_src, data.edge_type)), key=lambda x : (x[1], x[2]))
#def compute_e_to_distict_t(tu):
# num_edges = len(tu)
# all_node_distinct_types = 0
# cur_node = tu[0][1]
# type_set = set()
# type_set.add(tu[0][2])
# for i in range(1, len(tu)):
# if tu[i][1] == cur_node:
# type_set.add(tu[i][2])
# else:
# all_node_distinct_types += len(type_set)
# cur_node = tu[i][1]
# type_set.clear()
# type_set.add(tu[i][2])
# all_node_distinct_types += len(type_set)
# type_set.clear()
# #print('\n'.join([str(t) for t in tu]))
# print('num_edges:', num_edges, 'node distinct types', all_node_distinct_types)
# return num_edges/all_node_distinct_types
#r_forward = compute_e_to_distict_t(tu_forward)
#r_backward = compute_e_to_distict_t(tu_backward)
#print('ratio forward:', r_forward, 'ratio_backward:', r_backward)
# create model
model = EntityClassify(len(g),
args.n_hidden,
num_classes,
num_rels,
num_bases=args.num_bases,
num_hidden_layers=args.n_layers - 2,
dropout=args.dropout,
use_self_loop=args.use_self_loop,
use_cuda=use_cuda)
if use_cuda:
model.cuda()
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2norm)
# training loop
print("start training...")
forward_time = []
backward_time = []
model.train()
for epoch in range(args.n_epochs):
torch.cuda.synchronize()
optimizer.zero_grad()
t0 = time.time()
logits = model(g, feats, edge_type, edge_norm)
loss = F.cross_entropy(logits[train_idx], labels[train_idx])
torch.cuda.synchronize()
t1 = time.time()
loss.backward()
optimizer.step()
torch.cuda.synchronize()
t2 = time.time()
if epoch >= 3:
forward_time.append(t1 - t0)
backward_time.append(t2 - t1)
print("Epoch {:05d} | Train Forward Time(s) {:.4f} | Backward Time(s) {:.4f}".
format(epoch, forward_time[-1], backward_time[-1]))
train_acc = torch.sum(logits[train_idx].argmax(dim=1) == labels[train_idx]).item() / len(train_idx)
val_loss = F.cross_entropy(logits[val_idx], labels[val_idx])
val_acc = torch.sum(logits[val_idx].argmax(dim=1) == labels[val_idx]).item() / len(val_idx)
print("Train Accuracy: {:.4f} | Train Loss: {:.4f} | Validation Accuracy: {:.4f} | Validation loss: {:.4f}".
format(train_acc, loss.item(), val_acc, val_loss.item()))
print('max memory allocated', torch.cuda.max_memory_allocated())
model.eval()
logits = model.forward(g, feats, edge_type, edge_norm)
test_loss = F.cross_entropy(logits[test_idx], labels[test_idx])
test_acc = torch.sum(logits[test_idx].argmax(dim=1) == labels[test_idx]).item() / len(test_idx)
print("Test Accuracy: {:.4f} | Test loss: {:.4f}".format(test_acc, test_loss.item()))
print()
print("Mean forward time: {:4f}".format(np.mean(forward_time[len(forward_time) // 4:])))
print("Mean backward time: {:4f}".format(np.mean(backward_time[len(backward_time) // 4:])))
def main(args):
# graph
coo_adj = sp.load_npz("reddit_self_loop/reddit_self_loop_graph.npz")
graph = DGLGraph(coo_adj, readonly=True)
# features and labels
reddit_data = np.load("reddit_self_loop/reddit_data.npz")
features = reddit_data["feature"]
labels = reddit_data["label"]
num_labels = 41
# tarin/val/test indices
node_ids = reddit_data["node_ids"]
node_types = reddit_data["node_types"]
train_mask = (node_types == 1)
val_mask = (node_types == 2)
test_mask = (node_types == 3)
graph.ndata['train_mask'] = train_mask
graph.ndata['val_mask'] = val_mask
graph.ndata['test_mask'] = test_mask
graph.ndata['feat'] = features
graph.ndata['label'] = labels
features = torch.Tensor(features)
in_feats = features.shape[1]
labels = torch.LongTensor(labels)
train_nid = torch.LongTensor(np.where(train_mask == True)[0])
train_mask = torch.BoolTensor(train_mask)
val_nid = torch.LongTensor(np.where(val_mask == True)[0])
val_mask = torch.BoolTensor(val_mask)
test_nid = torch.LongTensor(np.where(test_mask == True)[0])
test_mask = torch.BoolTensor(test_mask)
g = dgl.graph(graph.all_edges()) # 转为HetroGraph
g.ndata['features'] = features
gpu = args.gpu
use_cuda = gpu >= 0 and torch.cuda.is_available()
if use_cuda:
torch.cuda.set_device(gpu)
g.to(torch.device('cuda:{}'.format(gpu)))
labels = labels.cuda()
fanouts = list(map(int, args.fan_out.split(',')))
sampler = Sample(g, fanouts)
# 将数据集打乱顺序，分多个batch，每个batch采样两个B
batch_size = args.batch_size
num_workers = args.num_workers
dataloader = DataLoader(dataset=train_nid.numpy(),
batch_size=batch_size,
collate_fn=sampler.obtain_Bs,
shuffle=True,
drop_last=False,
num_workers=num_workers)
# 设定模型
num_hid = args.num_hidden
ks = args.num_layers
dropout_r = args.dropout
agg = args.agg
bias = args.bias
norm = args.norm
model = GraphSAGE(in_feats,
num_hid,
num_labels,
ks,
bias=bias,
aggregator=agg,
activation=F.relu,
norm=norm,
dropout=dropout_r,
use_cuda=use_cuda)
if use_cuda:
model.cuda()
loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
l_r = args.lr
optimizer = torch.optim.Adam(model.parameters(), lr=l_r)
# acc
def compute_acc(pred, labels):
acc = (torch.argmax(pred, dim=1) == labels).float().sum() / len(pred)
return acc
# eval
def evaluation(model, g, labels, id, batch_size):
model.eval()
with torch.no_grad():
logits = model.infer(g, batch_size)
pred = logits[id]
label = labels[id]
model.train()
return compute_acc(pred, label)
# 训练、验证与测试
n_epochs = args.num_epochs
log_every = args.log_every
eval_every = args.eval_every
avg = 0
iter_tput = []
for epoch in range(n_epochs):
time_epoch_0 = time.time()
for step, blocks in enumerate(dataloader):
time_step = time.time()
input_nodes = blocks[0].srcdata[dgl.NID]
seeds = blocks[-1].dstdata[dgl.NID] # 最后一个block的dst为batch中的节点
batch_inputs = g.ndata['features'][input_nodes]
batch_labels = labels[seeds]
if use_cuda:
batch_inputs = batch_inputs.cuda()
batch_labels = batch_labels.cuda()
batch_pred = model(batch_inputs, blocks)
loss = loss_fcn(batch_pred, batch_labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
iter_tput.append(len(seeds) / (time.time() - time_step))
if step % log_every == 0:
acc = compute_acc(batch_pred, batch_labels)
print(
'Epoch {:05d} | Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f}'
.format(epoch, step, loss.item(), acc.item(),
np.mean(iter_tput[3:])))
time_epoch_1 = time.time()
print('Epoch Time(s): {:.4f}'.format(time_epoch_1 - time_epoch_0))
if epoch >= 5:
avg += time_epoch_1 - time_epoch_0
if epoch % eval_every == 0 and epoch != 0:
print('\n')
print('Eval-ing...')
time_ev_0 = time.time()
eval_acc = evaluation(model, g, labels, val_mask, batch_size)
time_ev_1 = time.time()
print('Eval Acc {:.4f} | Eval Time(s): {:.4f}'.format(
eval_acc, time_ev_1 - time_ev_0))
print('\n')
print('Eval-ing...')
time_ev_0 = time.time()
eval_acc = evaluation(model, g, labels, val_mask, batch_size)
time_ev_1 = time.time()
print('Eval Acc {:.4f} | Eval Time(s): {:.4f}'.format(
eval_acc, time_ev_1 - time_ev_0))
print('\n')
print('Testing...')
time_ev_0 = time.time()
test_acc = evaluation(model, g, labels, test_mask, batch_size)
time_ev_1 = time.time()
print('Test Acc {:.4f} | Eval Time(s): {:.4f}'.format(
test_acc, time_ev_1 - time_ev_0))
print('Finish!')
num_layers = 1
num_hidden = 16
infeat_dim = data.features.shape[1]
num_classes = data.num_labels
######################################################################
# Set up the DGL-PyTorch model and get the golden results
# -------------------------------------------------------
#
# The weights are trained with https://github.com/dmlc/dgl/blob/master/examples/pytorch/gcn/train.py
from tvm.contrib.download import download_testdata
from dgl import DGLGraph
features = torch.FloatTensor(data.features)
dgl_g = DGLGraph(g)
torch_model = GCN(dgl_g, infeat_dim, num_hidden, num_classes, num_layers,
F.relu)
# Download the pretrained weights
model_url = "https://homes.cs.washington.edu/~cyulin/media/gnn_model/gcn_%s.torch" % (
dataset)
model_path = download_testdata(model_url,
"gcn_%s.pickle" % (dataset),
module='gcn_model')
# Load the weights into the model
torch_model.load_state_dict(torch.load(model_path))
######################################################################
def load_data(data_filename, auxiliary_filename, active_thresh=1000, r=1):
# read the data file
# data_filename='covid_icd_for_gnn.csv'
raw_data = pd.read_csv(data_filename)
# Number of counties/ graph nodes
state_county = np.array(raw_data[['state', 'county']])
all_counties = np.unique(state_county[:, 0] + '_' + state_county[:, 1])
N = all_counties.size
# Number of time stamps
dates = np.unique(np.array(raw_data['date']))
T = dates.size
# static features - read from uszips
# auxiliary_filename = 'uszips.csv'
uszips = pd.read_csv(auxiliary_filename)
uszips_state_county = np.array(uszips[['state', 'county']])
uszips_all_counties = np.unique(uszips_state_county[:, 0] + '_' +
np.array(uszips_state_county)[:, 1])
all_popn = uszips.groupby(['state',
'county'])['population'].sum().reset_index()
all_density = uszips.groupby(['state',
'county'])['density'].sum().reset_index()
inter = raw_data.merge(all_popn)
raw_data_w_popn = inter.merge(all_density)
# Number of counties/ graph nodes
state_county_after_merge = np.array(raw_data_w_popn[['state', 'county']])
all_counties_after_merge = np.unique(state_county_after_merge[:, 0] + '_' +
state_county_after_merge[:, 1])
N_after_merge = all_counties_after_merge.size
# Number of time stamps
dates_after_merge = np.unique(np.array(raw_data_w_popn['date']))
T_after_merge = dates_after_merge.size
# note that some counties are exclusive to only one of the files in uszips.csv and covid-csv
# after merging them to consider county population, we retain only the common counties. This
# changes the number of counties to 2996 from 3004.
# duplicate rows, drop them
# 14546: FL.DESOTO
# 13777: DC
raw_data_w_popn = raw_data_w_popn.drop([14546, 13777])
raw_data_w_popn = raw_data_w_popn.drop(['country'], axis=1)
raw_data_w_popn['time_index'] = np.tile(np.arange(T_after_merge),
(N_after_merge))
data_for_sum = raw_data_w_popn.loc[:, ~raw_data_w_popn.columns.isin(
['latitude', 'longitude', 'density', 'county', 'time_index'])]
data_summed = data_for_sum.groupby(
['state', 'date']).sum().reset_index().drop(['Unnamed: 0'], axis=1)
data_for_mean = raw_data_w_popn[[
'state', 'date', 'latitude', 'longitude', 'density'
]]
data_mean = data_for_mean.groupby(['state', 'date']).mean().reset_index()
by_state = data_summed.merge(data_mean)
states_after_merge = np.unique(np.array(by_state['state']))
N_after_merge = states_after_merge.size
dates_after_merge = np.unique(np.array(by_state['date']))
T_after_merge = dates_after_merge.size
by_state_np = np.array(by_state)
# All features are in raw_data_w_popn
# smooth the data using the smooth1d function
win_len = 5
for n in range(N_after_merge):
by_state_np[n*T_after_merge:(n+1)*T_after_merge,4] = smooth1d(by_state_np[n*T_after_merge:(n+1)*T_after_merge,4]\
, win_len)
by_state_np[n*T_after_merge:(n+1)*T_after_merge,5] = smooth1d(by_state_np[n*T_after_merge:(n+1)*T_after_merge,5]\
, win_len)
info_for_edge = by_state[['latitude', 'longitude', 'population']]
popn = np.array(by_state['population']).astype('double')
by_state_np = by_state_np.drop(['population', 'n_bed'], axis=1)
dates = np.array(by_state)[:, 1]
active = np.array(by_state)[:, 2]
confirmed = np.array(by_state)[:, 3]
other_feat = np.array(
by_state_np)[:, 2:] # includes the active, confirmed # cases
active_cases = torch.from_numpy(
np.reshape(active, (N_after_merge, T_after_merge),
order='C').astype('float64'))
confirmed_cases = torch.from_numpy(
np.reshape(confirmed, (N_after_merge, T_after_merge),
order='C').astype('float64'))
# Reshape tensor into 3D array
feat_tensor = np.reshape(
np.array(other_feat),
(N_after_merge, T_after_merge, other_feat.shape[1]),
order='C')
feat_tensor = torch.from_numpy(feat_tensor.astype('float64'))
print("Feature tensor is of size ", feat_tensor.shape)
pop_data = torch.tensor(popn).view(N_after_merge, T_after_merge)
info_for_edge = np.reshape(np.array(info_for_edge),
(N_after_merge, T_after_merge, 3))
lat_long_pop = []
for ii in range(N_after_merge):
lat_long_pop.append(info_for_edge[ii][0])
W_mat = np.zeros((N_after_merge, N_after_merge))
for ii in range(N_after_merge):
lat1, lon1, pop1 = lat_long_pop[ii]
for jj in range(N_after_merge):
lat2, lon2, pop2 = lat_long_pop[jj]
if ii == jj:
W_mat[ii, jj] = 0
else:
W_mat[ii, jj] = calc_distance.gravity_law_commute_dist(
lat1, lon1, pop1, lat2, lon2, pop2, r)
# popn_density = np.reshape(np.array(by_state['density']),(N_after_merge,T_after_merge))
#
# W_mat = np.zeros((N_after_merge,N_after_merge))
#
# for ii in range(N_after_merge):
# pop_den1 = popn_density[ii,0]
# for jj in range(N_after_merge):
# pop_den2 = popn_density[jj,0]
# W_mat[ii,jj] = calc_distance.gravity_law_commute_dist(pop_den1, pop_den2)
# W_max = np.amax(W_mat)
#
# W_norm = W_mat/W_max
# W_norm_thresh = (W_norm > 1e-3)*W_norm
# W_sparse = sparse.coo_matrix(W_norm_thresh)
# W_sparse=sparse.coo_matrix(W_mat)
#
# values = W_sparse.data
# indices = np.vstack((W_sparse.row, W_sparse.col))
#
# i = torch.LongTensor(indices)
# v = torch.FloatTensor(values)
# shape = W_sparse.shape
#
# Adj=torch.sparse.FloatTensor(i, v, torch.Size(shape)).to_dense()
W_sparse = sparse.coo_matrix(W_mat)
adj = W_sparse + W_sparse.T.multiply(
W_sparse.T > W_sparse) - W_sparse.multiply(W_sparse.T > W_sparse)
def normalize_adj(mx):
"""Row-normalize sparse matrix"""
rowsum = np.array(mx.sum(1))
r_inv_sqrt = np.power(rowsum, -0.5).flatten()
r_inv_sqrt[np.isinf(r_inv_sqrt)] = 0.
r_mat_inv_sqrt = sparse.diags(r_inv_sqrt)
return mx.dot(r_mat_inv_sqrt).transpose().dot(r_mat_inv_sqrt)
adj = normalize_adj(adj + sparse.eye(adj.shape[0]))
Adj_norm = torch.FloatTensor(np.array(adj.todense()))
# row_sum=torch.sum(Adj,dim=1)
# row_sum[row_sum==0]=1
#
# invD_sqrt=torch.diag(torch.sqrt(row_sum.pow_(-1)))
# Adj_norm=torch.mm(invD_sqrt,(torch.mm(Adj,invD_sqrt)))
# Adj_norm=(Adj_norm-torch.diag(torch.diag(Adj_norm)))+torch.eye(Adj_norm.shape[0])
#
#
# W_sparse contains the adjacency matrix, which is good enough
# for the GCN model
# selected_counties = torch.where(torch.max(active_cases,dim=1)[0]>active_thresh)[0]
# sel=raw_data_w_popn.iloc[selected_counties*T_after_merge][['state','county']]
#
selected_states = torch.where(
torch.max(active_cases, dim=1)[0] > active_thresh)[0]
sel = by_state.iloc[selected_states * T_after_merge][['state']]
feat_tensor = feat_tensor[selected_states, :, :]
Adj_norm = Adj_norm[selected_states, :][:, selected_states]
confirmed_cases = confirmed_cases[selected_states, :]
active_cases = active_cases[selected_states, :]
pop_data = pop_data[selected_states, :]
# print("Using identity matrix as adjacency matrix")
# Adj_norm = torch.eye(Adj_norm.shape[0])
#
# g=DGLGraph()
# g.from_scipy_sparse_matrix(sparse.csr_matrix(Adj_norm.data.numpy()))
#
g = DGLGraph()
g.from_scipy_sparse_matrix(sparse.csr_matrix(Adj_norm.data.numpy()))
return feat_tensor, Adj_norm, active_cases, confirmed_cases, pop_data, sel, g
def main(args):
# load and preprocess dataset
data = load_data(args)
if args.self_loop and not args.dataset.startswith('reddit'):
data.graph.add_edges_from([(i,i) for i in range(len(data.graph))])
train_nid = np.nonzero(data.train_mask)[0].astype(np.int64)
test_nid = np.nonzero(data.test_mask)[0].astype(np.int64)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
n_train_samples = train_mask.sum().item()
n_val_samples = val_mask.sum().item()
n_test_samples = test_mask.sum().item()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes,
n_train_samples,
n_val_samples,
n_test_samples))
# create GCN model
g = DGLGraph(data.graph, readonly=True)
norm = 1. / g.in_degrees().float().unsqueeze(1)
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
norm = norm.cuda()
g.ndata['features'] = features
num_neighbors = args.num_neighbors
g.ndata['norm'] = norm
model = GCNSampling(in_feats,
args.n_hidden,
n_classes,
args.n_layers,
F.relu,
args.dropout)
if cuda:
model.cuda()
loss_fcn = nn.CrossEntropyLoss()
infer_model = GCNInfer(in_feats,
args.n_hidden,
n_classes,
args.n_layers,
F.relu)
if cuda:
infer_model.cuda()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
for epoch in range(args.n_epochs):
for nf in dgl.contrib.sampling.NeighborSampler(g, args.batch_size,
args.num_neighbors,
neighbor_type='in',
shuffle=True,
num_workers=32,
num_hops=args.n_layers+1,
seed_nodes=train_nid):
nf.copy_from_parent()
model.train()
# forward
pred = model(nf)
batch_nids = nf.layer_parent_nid(-1).to(device=pred.device, dtype=torch.long)
batch_labels = labels[batch_nids]
loss = loss_fcn(pred, batch_labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
for infer_param, param in zip(infer_model.parameters(), model.parameters()):
infer_param.data.copy_(param.data)
num_acc = 0.
for nf in dgl.contrib.sampling.NeighborSampler(g, args.test_batch_size,
g.number_of_nodes(),
neighbor_type='in',
num_workers=32,
num_hops=args.n_layers+1,
seed_nodes=test_nid):
nf.copy_from_parent()
infer_model.eval()
with torch.no_grad():
pred = infer_model(nf)
batch_nids = nf.layer_parent_nid(-1).to(device=pred.device, dtype=torch.long)
batch_labels = labels[batch_nids]
num_acc += (pred.argmax(dim=1) == batch_labels).sum().cpu().item()
print("Test Accuracy {:.4f}". format(num_acc/n_test_samples))
def trainer(rank, world_size, args, backend='nccl'):
# init multi process
init_process(rank, world_size, backend)
# load data
dataname = os.path.basename(args.dataset)
remote_g = dgl.contrib.graph_store.create_graph_from_store(
dataname, "shared_mem")
adj, t2fid = data.get_sub_train_graph(args.dataset, rank, world_size)
g = DGLGraph(adj, readonly=True)
n_classes = args.n_classes
train_nid = data.get_sub_train_nid(args.dataset, rank, world_size)
sub_labels = data.get_sub_train_labels(args.dataset, rank, world_size)
labels = np.zeros(np.max(train_nid) + 1, dtype=np.int)
labels[train_nid] = sub_labels
# to torch tensor
t2fid = torch.LongTensor(t2fid)
labels = torch.LongTensor(labels)
if args.preprocess:
embed_names = ['features', 'neigh']
else:
embed_names = ['features']
cacher = storage.GraphCacheServer(remote_g, adj.shape[0], t2fid, rank)
cacher.init_field(embed_names)
cacher.log = False
# prepare model
num_hops = args.n_layers if args.preprocess else args.n_layers + 1
model = GraphSageSampling(args.feat_size, args.n_hidden, n_classes,
args.n_layers, F.relu, args.dropout, 'mean',
args.preprocess)
loss_fcn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
model.cuda(rank)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])
ctx = torch.device(rank)
if args.remote_sample:
sampler = SampleLoader(g, rank, one2all=False)
else:
sampler = dgl.contrib.sampling.NeighborSampler(
g,
args.batch_size,
args.num_neighbors,
neighbor_type='in',
shuffle=True,
num_workers=args.num_workers,
num_hops=num_hops,
seed_nodes=train_nid,
prefetch=True)
# start training
epoch_dur = []
tic = time.time()
with torch.autograd.profiler.profile(enabled=(rank == 0),
use_cuda=True) as prof:
for epoch in range(args.n_epochs):
model.train()
epoch_start_time = time.time()
step = 0
for nf in sampler:
with torch.autograd.profiler.record_function('gpu-load'):
cacher.fetch_data(nf)
batch_nids = nf.layer_parent_nid(-1)
label = labels[batch_nids]
label = label.cuda(rank, non_blocking=True)
with torch.autograd.profiler.record_function('gpu-compute'):
pred = model(nf)
loss = loss_fcn(pred, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
step += 1
if epoch == 0 and step == 1:
cacher.auto_cache(g, embed_names)
if rank == 0 and step % 20 == 0:
print('epoch [{}] step [{}]. Loss: {:.4f}'.format(
epoch + 1, step, loss.item()))
if rank == 0:
epoch_dur.append(time.time() - epoch_start_time)
print('Epoch average time: {:.4f}'.format(
np.mean(np.array(epoch_dur[2:]))))
if cacher.log:
miss_rate = cacher.get_miss_rate()
print('Epoch average miss rate: {:.4f}'.format(miss_rate))
toc = time.time()
if rank == 0:
print(prof.key_averages().table(sort_by='cuda_time_total'))
print('Total Time: {:.4f}s'.format(toc - tic))
def main(args):
torch.manual_seed(args.rnd_seed)
np.random.seed(args.rnd_seed)
random.seed(args.rnd_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
multitask_data = set([
'ppi', 'amazon', 'amazon-0.1', 'amazon-0.3', 'amazon2M', 'amazon2M-47'
])
multitask = args.dataset in multitask_data
# load and preprocess dataset
data = load_data(args)
train_nid = np.nonzero(data.train_mask)[0].astype(np.int64)
test_nid = np.nonzero(data.test_mask)[0].astype(np.int64)
# Normalize features
if args.normalize:
train_feats = data.features[train_nid]
scaler = sklearn.preprocessing.StandardScaler()
scaler.fit(train_feats)
features = scaler.transform(data.features)
else:
features = data.features
features = torch.FloatTensor(features)
if not multitask:
labels = torch.LongTensor(data.labels)
else:
labels = torch.FloatTensor(data.labels)
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
n_train_samples = train_mask.sum().item()
n_val_samples = val_mask.sum().item()
n_test_samples = test_mask.sum().item()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes, n_train_samples, n_val_samples, n_test_samples))
# create GCN model
g = data.graph
if args.self_loop and not args.dataset.startswith('reddit'):
g.remove_edges_from(g.selfloop_edges())
g.add_edges_from(zip(g.nodes(), g.nodes()))
print("adding self-loop edges")
g = DGLGraph(g, readonly=True)
# set device for dataset tensors
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
print(torch.cuda.get_device_name(0))
g.ndata['features'] = features
g.ndata['labels'] = labels
g.ndata['train_mask'] = train_mask
print('labels shape:', labels.shape)
cluster_iterator = ClusterIter(args.dataset,
g,
args.psize,
args.batch_size,
train_nid,
use_pp=args.use_pp)
print("features shape, ", features.shape)
model_sel = {'GCN': GCNCluster, 'graphsage': GraphSAGE}
model_class = model_sel[args.model_type]
print('using model:', model_class)
model = model_class(in_feats, args.n_hidden, n_classes, args.n_layers,
F.relu, args.dropout, args.use_pp)
if cuda:
model.cuda()
# logger and so on
#log_dir = save_log_dir(args)
#writer = SummaryWriter(log_dir)
#logger = Logger(os.path.join(log_dir, 'loggings'))
#logger.write(args)
# Loss function
if multitask:
print('Using multi-label loss')
loss_f = nn.BCEWithLogitsLoss()
else:
print('Using multi-class loss')
loss_f = nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
# set train_nids to cuda tensor
if cuda:
train_nid = torch.from_numpy(train_nid).cuda()
print("current memory after model before training",
torch.cuda.memory_allocated(device=train_nid.device) / 1024 / 1024)
start_time = time.time()
best_f1 = -1
for epoch in range(args.n_epochs):
for j, cluster in enumerate(cluster_iterator):
# sync with upper level training graph
cluster.copy_from_parent()
model.train()
# forward
pred = model(cluster)
batch_labels = cluster.ndata['labels']
batch_train_mask = cluster.ndata['train_mask']
loss = loss_f(pred[batch_train_mask],
batch_labels[batch_train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
# in PPI case, `log_every` is chosen to log one time per epoch.
# Choose your log freq dynamically when you want more info within one epoch
if j % args.log_every == 0:
print(
f"epoch:{epoch}/{args.n_epochs}, Iteration {j}/{len(cluster_iterator)}:training loss",
loss.item())
#writer.add_scalar('train/loss', loss.item(),
# global_step=j + epoch * len(cluster_iterator))
print("current memory:",
torch.cuda.memory_allocated(device=pred.device) / 1024 / 1024)
# evaluate
if epoch % args.val_every == 0:
val_f1_mic, val_f1_mac = evaluate(model, g, labels, val_mask,
multitask)
print("Val F1-mic{:.4f}, Val F1-mac{:.4f}".format(
val_f1_mic, val_f1_mac))
if val_f1_mic > best_f1:
best_f1 = val_f1_mic
print('new best val f1:', best_f1)
torch.save(model.state_dict(),
os.path.join(log_dir, 'best_model.pkl'))
#writer.add_scalar('val/f1-mic', val_f1_mic, global_step=epoch)
#writer.add_scalar('val/f1-mac', val_f1_mac, global_step=epoch)
end_time = time.time()
print(f'training using time {start_time-end_time}')
# test
if args.use_val:
model.load_state_dict(
torch.load(os.path.join(log_dir, 'best_model.pkl')))
test_f1_mic, test_f1_mac = evaluate(model, g, labels, test_mask, multitask)
print("Test F1-mic{:.4f}, Test F1-mac{:.4f}".format(
test_f1_mic, test_f1_mac))
def main(args):
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.FloatTensor(data.labels)
train_mask = torch.BoolTensor(data.train_mask)
val_mask = torch.BoolTensor(data.val_mask)
test_mask = torch.BoolTensor(data.test_mask)
g = data.graph
n_feats = features.shape[1]
n_labels = data.num_labels
n_edges = g.number_of_edges()
print("""----Data statistics------'
#Features %d
#Edges %d
#Labels %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_feats, n_edges, n_labels, train_mask.int().sum().item(),
val_mask.int().sum().item(), test_mask.int().sum().item()))
dataset_train = CampusDataset(features, labels)
dict_users = iid_users(dataset_train, args.n_users)
if args.gnnbase == 'gcn':
g = DGLGraph(g)
n_edges = g.number_of_edges()
degs = g.in_degrees().float()
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
g.ndata['norm'] = norm.unsqueeze(1)
model = GCN(g, n_feats, args.n_hidden, n_labels, args.n_layers, F.relu,
args.dropout)
if args.gnnbase == 'gat':
g.remove_edges_from(nx.selfloop_edges(g))
g = DGLGraph(g)
g.add_edges(g.nodes(), g.nodes())
n_edges = g.number_of_edges()
heads = ([args.n_heads] * args.n_layers) + [args.n_out_heads]
model = GAT(g, args.n_layers, n_feats, args.n_hidden, n_labels, heads,
F.elu, args.in_drop, args.attn_drop, args.negative_slope,
args.residual)
if args.gnnbase == 'sage':
g.remove_edges_from(nx.selfloop_edges(g))
g = DGLGraph(g)
n_edges = g.number_of_edges()
model = GraphSAGE(g, n_feats, args.n_hidden, n_labels, args.n_layers,
F.relu, args.dropout, args.aggregator_type)
print(model)
model.train()
w_glob = model.state_dict()
loss_train = []
timecost = []
for epoch in range(args.n_epochs):
time_begin = time.time()
w_locals, loss_locals = [], []
m = max(int(args.frac * args.n_users), 1)
idxs_users = np.random.choice(range(args.n_users), m, replace=False)
for idx in idxs_users:
local = LocalUpdate(args=args,
dataset=dataset_train,
idxs=dict_users[idx],
mask=train_mask)
w, loss = local.train(model=copy.deepcopy(model))
w_locals.append(copy.deepcopy(w))
loss_locals.append(copy.deepcopy(loss))
w_glob = FedAvg(w_locals)
model.load_state_dict(w_glob)
time_end = time.time()
timecost.append(time_end - time_begin)
loss_avg = sum(loss_locals) / len(loss_locals)
print('Epoch {:3d}, Average loss {:.3f}'.format(epoch, loss_avg))
loss_train.append(loss_avg)
train_errX, train_errY = eval_error(model, features, labels,
train_mask)
val_errX, val_errY = eval_error(model, features, labels, val_mask)
test_errX, test_errY = eval_error(model, features, labels, test_mask)
print(
"Epoch {:3d} | TrainRMSEX {:.4f} | TrainRMSEY {:.4f} | ValRMSEX {:.4f} | ValRMSEY {:.4f} | TestRMSEX {:.4f} | TestRMSEY {:.4f}"
.format(epoch, train_errX, train_errY, val_errX, val_errY,
test_errX, test_errY))
print("Time cost {:.4f}".format(sum(timecost) / args.n_epochs))
base_errX, base_errY = calc_error(features[test_mask, :2],
labels[test_mask])
print("TestRMSEX-Base {:.4f} | TestRMSEY-Base {:.4f}".format(
base_errX, base_errY))
def main(args):
# load and preprocess dataset
data = load_data(args)
features = mx.nd.array(data.features)
labels = mx.nd.array(data.labels)
train_mask = mx.nd.array(data.train_mask)
val_mask = mx.nd.array(data.val_mask)
test_mask = mx.nd.array(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes, train_mask.sum().asscalar(),
val_mask.sum().asscalar(), test_mask.sum().asscalar()))
if args.gpu < 0:
cuda = False
ctx = mx.cpu(0)
else:
cuda = True
ctx = mx.gpu(args.gpu)
features = features.as_in_context(ctx)
labels = labels.as_in_context(ctx)
train_mask = train_mask.as_in_context(ctx)
val_mask = val_mask.as_in_context(ctx)
test_mask = test_mask.as_in_context(ctx)
# create GCN model
g = data.graph
if args.self_loop:
g.remove_edges_from(nx.selfloop_edges(g))
g.add_edges_from(zip(g.nodes(), g.nodes()))
g = DGLGraph(g)
# normalization
degs = g.in_degrees().astype('float32')
norm = mx.nd.power(degs, -0.5)
if cuda:
norm = norm.as_in_context(ctx)
g.ndata['norm'] = mx.nd.expand_dims(norm, 1)
model = GCN(g, in_feats, args.n_hidden, n_classes, args.n_layers,
mx.nd.relu, args.dropout)
model.initialize(ctx=ctx)
n_train_samples = train_mask.sum().asscalar()
loss_fcn = gluon.loss.SoftmaxCELoss()
# use optimizer
print(model.collect_params())
trainer = gluon.Trainer(model.collect_params(), 'adam', {
'learning_rate': args.lr,
'wd': args.weight_decay
})
# initialize graph
dur = []
for epoch in range(args.n_epochs):
if epoch >= 3:
t0 = time.time()
# forward
with mx.autograd.record():
pred = model(features)
loss = loss_fcn(pred, labels, mx.nd.expand_dims(train_mask, 1))
loss = loss.sum() / n_train_samples
loss.backward()
trainer.step(batch_size=1)
if epoch >= 3:
loss.asscalar()
dur.append(time.time() - t0)
acc = evaluate(model, features, labels, val_mask)
print(
"Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
"ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur),
loss.asscalar(), acc,
n_edges / np.mean(dur) / 1000))
# test set accuracy
acc = evaluate(model, features, labels, test_mask)
print("Test accuracy {:.2%}".format(acc))
def main(args):
# load and preprocess dataset
data = load_data(args)
if args.gpu >= 0:
ctx = mx.gpu(args.gpu)
else:
ctx = mx.cpu()
if args.self_loop and not args.dataset.startswith('reddit'):
data.graph.add_edges_from([(i, i) for i in range(len(data.graph))])
train_nid = mx.nd.array(np.nonzero(data.train_mask)[0]).astype(
np.int64).as_in_context(ctx)
test_nid = mx.nd.array(np.nonzero(data.test_mask)[0]).astype(
np.int64).as_in_context(ctx)
features = mx.nd.array(data.features).as_in_context(ctx)
labels = mx.nd.array(data.labels).as_in_context(ctx)
train_mask = mx.nd.array(data.train_mask).as_in_context(ctx)
val_mask = mx.nd.array(data.val_mask).as_in_context(ctx)
test_mask = mx.nd.array(data.test_mask).as_in_context(ctx)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
n_train_samples = train_mask.sum().asscalar()
n_val_samples = val_mask.sum().asscalar()
n_test_samples = test_mask.sum().asscalar()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes, n_train_samples, n_val_samples, n_test_samples))
# create GCN model
g = DGLGraph(data.graph, readonly=True)
g.ndata['features'] = features
num_neighbors = args.num_neighbors
degs = g.in_degrees().astype('float32').as_in_context(ctx)
norm = mx.nd.expand_dims(1. / degs, 1)
g.ndata['norm'] = norm
model = GCNSampling(in_feats,
args.n_hidden,
n_classes,
args.n_layers,
mx.nd.relu,
args.dropout,
prefix='GCN')
model.initialize(ctx=ctx)
loss_fcn = gluon.loss.SoftmaxCELoss()
infer_model = GCNInfer(in_feats,
args.n_hidden,
n_classes,
args.n_layers,
mx.nd.relu,
prefix='GCN')
infer_model.initialize(ctx=ctx)
# use optimizer
print(model.collect_params())
trainer = gluon.Trainer(model.collect_params(),
'adam', {
'learning_rate': args.lr,
'wd': args.weight_decay
},
kvstore=mx.kv.create('local'))
# initialize graph
dur = []
for epoch in range(args.n_epochs):
index = 0
for nf in dgl.contrib.sampling.NeighborSampler(g,
args.batch_size,
args.num_neighbors,
neighbor_type='in',
shuffle=True,
num_hops=args.n_layers +
1,
seed_nodes=train_nid):
print(index)
index = index + 1
nf.copy_from_parent()
# forward
with mx.autograd.record():
pred = model(nf)
batch_nids = nf.layer_parent_nid(-1).astype(
'int64').as_in_context(ctx)
batch_labels = labels[batch_nids]
loss = loss_fcn(pred, batch_labels)
loss = loss.sum() / len(batch_nids)
loss.backward()
trainer.step(batch_size=1)
infer_params = infer_model.collect_params()
for key in infer_params:
idx = trainer._param2idx[key]
trainer._kvstore.pull(idx, out=infer_params[key].data())
num_acc = 0.
for nf in dgl.contrib.sampling.NeighborSampler(g,
args.test_batch_size,
g.number_of_nodes(),
neighbor_type='in',
num_hops=args.n_layers +
1,
seed_nodes=test_nid):
nf.copy_from_parent()
pred = infer_model(nf)
batch_nids = nf.layer_parent_nid(-1).astype('int64').as_in_context(
ctx)
batch_labels = labels[batch_nids]
num_acc += (pred.argmax(axis=1) == batch_labels).sum().asscalar()
print("Test Accuracy {:.4f}".format(num_acc / n_test_samples))
def main(args):
# load and preprocess dataset
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
if hasattr(torch, 'BoolTensor'):
train_mask = torch.BoolTensor(data.train_mask)
val_mask = torch.BoolTensor(data.val_mask)
test_mask = torch.BoolTensor(data.test_mask)
else:
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes, train_mask.int().sum().item(),
val_mask.int().sum().item(), test_mask.int().sum().item()))
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
# graph preprocess and calculate normalization factor
g = data.graph
# add self loop
if args.self_loop:
g.remove_edges_from(nx.selfloop_edges(g))
g.add_edges_from(zip(g.nodes(), g.nodes()))
g = DGLGraph(g)
n_edges = g.number_of_edges()
# normalization
degs = g.in_degrees().float()
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
if cuda:
norm = norm.cuda()
g.ndata['norm'] = norm.unsqueeze(1)
# create GCN model
model = GCN(g, in_feats, args.n_hidden, n_classes, args.n_layers, F.relu,
args.dropout)
if cuda:
model.cuda()
loss_fcn = torch.nn.CrossEntropyLoss()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
for epoch in range(args.n_epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(features)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch >= 3:
dur.append(time.time() - t0)
acc = evaluate(model, features, labels, val_mask)
print(
"Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
"ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item(),
acc, n_edges / np.mean(dur) / 1000))
print()
acc = evaluate(model, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc))
def test_nx_conversion():
# check conversion between networkx and DGLGraph
def _check_nx_feature(nxg, nf, ef):
# check node and edge feature of nxg
# this is used to check to_networkx
num_nodes = len(nxg)
num_edges = nxg.size()
if num_nodes > 0:
node_feat = ddict(list)
for nid, attr in nxg.nodes(data=True):
assert len(attr) == len(nf)
for k in nxg.nodes[nid]:
node_feat[k].append(F.unsqueeze(attr[k], 0))
for k in node_feat:
feat = F.cat(node_feat[k], 0)
assert F.allclose(feat, nf[k])
else:
assert len(nf) == 0
if num_edges > 0:
edge_feat = ddict(lambda: [0] * num_edges)
for u, v, attr in nxg.edges(data=True):
assert len(attr) == len(ef) + 1 # extra id
eid = attr['id']
for k in ef:
edge_feat[k][eid] = F.unsqueeze(attr[k], 0)
for k in edge_feat:
feat = F.cat(edge_feat[k], 0)
assert F.allclose(feat, ef[k])
else:
assert len(ef) == 0
n1 = F.randn((5, 3))
n2 = F.randn((5, 10))
n3 = F.randn((5, 4))
e1 = F.randn((4, 5))
e2 = F.randn((4, 7))
g = DGLGraph(multigraph=True)
g.add_nodes(5)
g.add_edges([0,1,3,4], [2,4,0,3])
g.ndata.update({'n1': n1, 'n2': n2, 'n3': n3})
g.edata.update({'e1': e1, 'e2': e2})
# convert to networkx
nxg = g.to_networkx(node_attrs=['n1', 'n3'], edge_attrs=['e1', 'e2'])
assert len(nxg) == 5
assert nxg.size() == 4
_check_nx_feature(nxg, {'n1': n1, 'n3': n3}, {'e1': e1, 'e2': e2})
# convert to DGLGraph, nx graph has id in edge feature
# use id feature to test non-tensor copy
g.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'])
# check graph size
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
# check number of features
# test with existing dglgraph (so existing features should be cleared)
assert len(g.ndata) == 1
assert len(g.edata) == 2
# check feature values
assert F.allclose(g.ndata['n1'], n1)
# with id in nx edge feature, e1 should follow original order
assert F.allclose(g.edata['e1'], e1)
assert F.array_equal(g.get_e_repr()['id'], F.arange(0, 4))
# test conversion after modifying DGLGraph
g.pop_e_repr('id') # pop id so we don't need to provide id when adding edges
new_n = F.randn((2, 3))
new_e = F.randn((3, 5))
g.add_nodes(2, data={'n1': new_n})
# add three edges, one is a multi-edge
g.add_edges([3, 6, 0], [4, 5, 2], data={'e1': new_e})
n1 = F.cat((n1, new_n), 0)
e1 = F.cat((e1, new_e), 0)
# convert to networkx again
nxg = g.to_networkx(node_attrs=['n1'], edge_attrs=['e1'])
assert len(nxg) == 7
assert nxg.size() == 7
_check_nx_feature(nxg, {'n1': n1}, {'e1': e1})
# now test convert from networkx without id in edge feature
# first pop id in edge feature
for _, _, attr in nxg.edges(data=True):
attr.pop('id')
# test with a new graph
g = DGLGraph(multigraph=True)
g.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1'])
# check graph size
assert g.number_of_nodes() == 7
assert g.number_of_edges() == 7
# check number of features
assert len(g.ndata) == 1
assert len(g.edata) == 1
# check feature values
assert F.allclose(g.ndata['n1'], n1)
# edge feature order follows nxg.edges()
edge_feat = []
for _, _, attr in nxg.edges(data=True):
edge_feat.append(F.unsqueeze(attr['e1'], 0))
edge_feat = F.cat(edge_feat, 0)
assert F.allclose(g.edata['e1'], edge_feat)
import matplotlib
matplotlib.use('Agg') # Must be before importing matplotlib.pyplot or pylab!
import matplotlib.pyplot as plt
from optparse import OptionParser
from models.vanilla_gcn import VanillaGCN
from models.lstm_gcn import LSTMGCN
import seaborn as sns
sns.set()
import pdb
import config
path = os.path.join(config.GRAPH_DIR, "c17_10e4.graphml")
print("Loading train graph from {}".format(path))
graph_train = nx.read_graphml(path)
g_train = DGLGraph(graph_train)
all_types = []
all_ids_train = list(graph_train.nodes.keys())
for node_id in graph_train.nodes:
if graph_train.nodes[node_id]['gtype'] not in all_types:
all_types.append(graph_train.nodes[node_id]['gtype'])
print(all_types)
features_list = [x.strip() for x in options.features.split(",")
] if options.features else []
feature_dim = 1 #1 is the default feature_dim when no features are provided
for feature_name in features_list:
if feature_name == "gtype":
feature_dim += len(all_types)
else:
def main(args):
coo_adj, feat = data.get_graph_data(args.dataset)
graph = dgl.DGLGraph(coo_adj, readonly=True)
features = torch.FloatTensor(feat)
graph_name = os.path.basename(args.dataset)
vnum = graph.number_of_nodes()
enum = graph.number_of_edges()
feat_size = feat.shape[1]
print('=' * 30)
print("Graph Name: {}\nNodes Num: {}\tEdges Num: {}\nFeature Size: {}"
.format(graph_name, vnum, enum, feat_size)
)
print('=' * 30)
# create server
g = dgl.contrib.graph_store.create_graph_store_server(
graph, graph_name,
'shared_mem', args.num_workers,
False, edge_dir='in')
# calculate norm for gcn
dgl_g = DGLGraph(graph, readonly=True)
if args.model == 'gcn':
dgl_g = DGLGraph(graph, readonly=True)
norm = 1. / dgl_g.in_degrees().float().unsqueeze(1)
# preprocess
if args.preprocess:
print('Preprocessing features...')
dgl_g.ndata['norm'] = norm
dgl_g.ndata['features'] = features
dgl_g.update_all(fn.copy_src(src='features', out='m'),
fn.sum(msg='m', out='preprocess'),
lambda node : {'preprocess': node.data['preprocess'] * node.data['norm']})
features = dgl_g.ndata['preprocess']
g.ndata['norm'] = norm
g.ndata['features'] = features
del dgl_g
elif args.model == 'graphsage':
if args.preprocess: # for simple preprocessing
print('preprocessing: warning: jusy copy')
g.ndata['neigh'] = features
g.ndata['features'] = features
# remote sampler
if args.sample:
subgraph = []
sub_trainnid = []
for rank in range(args.num_workers):
subadj, _ = data.get_sub_train_graph(args.dataset, rank, args.num_workers)
train_nid = data.get_sub_train_nid(args.dataset, rank, args.num_workers)
subgraph.append(dgl.DGLGraph(subadj, readonly=True))
sub_trainnid.append(train_nid)
hops = args.gnn_layers - 1 if args.preprocess else args.gnn_layers
print('Expected trainer#: {}. Start sampling at server end...'.format(args.num_workers))
deliver = SampleDeliver(subgraph, sub_trainnid, args.num_neighbors, hops, args.num_workers)
deliver.async_sample(args.n_epochs, args.batch_size, one2all=args.one2all)
print('start running graph server on dataset: {}'.format(graph_name))
g.run()