1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
|
from itertools import product
import dgl
import torch
from dgl import DGLGraph
from dgl.contrib.data import load_data
from dgl.data import citation_graph
from runtime.dgl.gat import GAT, GATSPMV
from runtime.dgl.gcn import GCN, GCNSPMV
from runtime.dgl.hidden import HiddenPrint
from runtime.dgl.rgcn import RGCN, RGCNSPMV
from runtime.dgl.train import train_runtime
if torch.cuda.is_available():
device = torch.device('cuda')
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = torch.device('mps')
else:
device = torch.device('cpu')
with HiddenPrint():
Cora = citation_graph.load_cora()
CiteSeer = citation_graph.load_citeseer()
PubMed = citation_graph.load_pubmed()
MUTAG = load_data('mutag') # fair comparison
# One training run before we start tracking duration to warm up GPU.
g = DGLGraph(Cora.graph)
g.set_n_initializer(dgl.init.zero_initializer)
g.add_edges(g.nodes(), g.nodes())
norm = torch.pow(g.in_degrees().float(), -0.5)
norm[torch.isinf(norm)] = 0
g.ndata['norm'] = norm.unsqueeze(1).to(device)
model = GCNSPMV(g, Cora.features.shape[1], Cora.num_labels).to(device)
train_runtime(model, Cora, epochs=200, device=device)
for d, Net in product([Cora, CiteSeer, PubMed], [GCN, GCNSPMV, GAT, GATSPMV]):
g = DGLGraph(d.graph)
g.set_n_initializer(dgl.init.zero_initializer)
g.add_edges(g.nodes(), g.nodes())
norm = torch.pow(g.in_degrees().float(), -0.5)
norm[torch.isinf(norm)] = 0
g.ndata['norm'] = norm.unsqueeze(1).to(device)
model = Net(g, d.features.shape[1], d.num_labels).to(device)
t = train_runtime(model, d, epochs=200, device=device)
print(f'{d.name} - {Net.__name__}: {t:.2f}s')
for d, Net in product([MUTAG], [RGCN, RGCNSPMV]):
g = DGLGraph()
g.add_nodes(d.num_nodes)
g.add_edges(d.edge_src, d.edge_dst)
edge_type = torch.from_numpy(d.edge_type).to(device)
edge_norm = torch.from_numpy(d.edge_norm).to(device)
g.edata.update({'type': edge_type, 'norm': edge_norm})
g.ndata['id'] = torch.arange(d.num_nodes, dtype=torch.long, device=device)
model = Net(g, d.num_nodes, d.num_classes, d.num_rels)
t = train_runtime(model, d, epochs=200, device=device)
print(f'{d.name} - {Net.__name__}: {t:.2f}s')
|
