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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
|
import argparse
import os.path as osp
import torch
import torch.nn.functional as F
from ogb.graphproppred import Evaluator
from ogb.graphproppred import PygGraphPropPredDataset as OGBG
from ogb.graphproppred.mol_encoder import AtomEncoder
from torch.nn import BatchNorm1d, Linear, ReLU, Sequential
from torch.optim.lr_scheduler import ReduceLROnPlateau
import torch_geometric.transforms as T
from torch_geometric.loader import DataLoader
from torch_geometric.nn import EGConv, global_mean_pool
from torch_geometric.typing import WITH_TORCH_SPARSE
if not WITH_TORCH_SPARSE:
quit("This example requires 'torch-sparse'")
parser = argparse.ArgumentParser()
parser.add_argument('--use_multi_aggregators', action='store_true',
help='Switch between EGC-S and EGC-M')
args = parser.parse_args()
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'OGB')
dataset = OGBG('ogbg-molhiv', path, pre_transform=T.ToSparseTensor())
evaluator = Evaluator('ogbg-molhiv')
split_idx = dataset.get_idx_split()
train_dataset = dataset[split_idx['train']]
val_dataset = dataset[split_idx['valid']]
test_dataset = dataset[split_idx['test']]
train_loader = DataLoader(train_dataset, batch_size=32, num_workers=4,
shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=256)
test_loader = DataLoader(test_dataset, batch_size=256)
class Net(torch.nn.Module):
def __init__(self, hidden_channels, num_layers, num_heads, num_bases):
super().__init__()
if args.use_multi_aggregators:
aggregators = ['sum', 'mean', 'max']
else:
aggregators = ['symnorm']
self.encoder = AtomEncoder(hidden_channels)
self.convs = torch.nn.ModuleList()
self.norms = torch.nn.ModuleList()
for _ in range(num_layers):
self.convs.append(
EGConv(hidden_channels, hidden_channels, aggregators,
num_heads, num_bases))
self.norms.append(BatchNorm1d(hidden_channels))
self.mlp = Sequential(
Linear(hidden_channels, hidden_channels // 2, bias=False),
BatchNorm1d(hidden_channels // 2),
ReLU(inplace=True),
Linear(hidden_channels // 2, hidden_channels // 4, bias=False),
BatchNorm1d(hidden_channels // 4),
ReLU(inplace=True),
Linear(hidden_channels // 4, 1),
)
def forward(self, x, adj_t, batch):
adj_t = adj_t.set_value(None) # EGConv works without any edge features
x = self.encoder(x)
for conv, norm in zip(self.convs, self.norms):
h = conv(x, adj_t)
h = norm(h)
h = h.relu_()
x = x + h
x = global_mean_pool(x, batch)
return self.mlp(x)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net(hidden_channels=236, num_layers=4, num_heads=4,
num_bases=4).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=20,
min_lr=1e-5)
def train():
model.train()
total_loss = total_examples = 0
for data in train_loader:
data = data.to(device)
optimizer.zero_grad()
out = model(data.x, data.adj_t, data.batch)
loss = F.binary_cross_entropy_with_logits(out, data.y.to(torch.float))
loss.backward()
optimizer.step()
total_loss += float(loss) * data.num_graphs
total_examples += data.num_graphs
return total_loss / total_examples
@torch.no_grad()
def evaluate(loader):
model.eval()
y_pred, y_true = [], []
for data in loader:
data = data.to(device)
pred = model(data.x, data.adj_t, data.batch)
y_pred.append(pred.cpu())
y_true.append(data.y.cpu())
y_true = torch.cat(y_true, dim=0)
y_pred = torch.cat(y_pred, dim=0)
return evaluator.eval({'y_true': y_true, 'y_pred': y_pred})['rocauc']
for epoch in range(1, 31):
loss = train()
val_rocauc = evaluate(val_loader)
test_rocauc = evaluate(test_loader)
scheduler.step(val_rocauc)
print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Val: {val_rocauc:.4f}, '
f'Test: {test_rocauc:.4f}')
|
