-
Notifications
You must be signed in to change notification settings - Fork 3
/
train_large.py
136 lines (103 loc) · 4.9 KB
/
train_large.py
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
135
136
# Description: This file is used to train GNN models on large datasets. The trained model is saved
# in the pretrained folder. The model is trained using NeighborLoader, which is not supported
# in train.py.
# The trained model is used for benchmarking explanation methods.
import argparse
import copy
import os
import sys
import torch
import torch.nn.functional as F
from torch_geometric.loader import NeighborLoader
from tqdm import tqdm
from dataset.utils import get_model_data_config
from gnnshap.utils import pruned_comp_graph
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', default='Reddit', type=str)
args = parser.parse_args()
dataset_name = args.dataset
def train(epoch):
model.train()
pbar = tqdm(total=int(len(train_loader.dataset)))
pbar.set_description(f'Epoch {epoch:02d}')
total_loss = total_correct = total_examples = 0
for batch in train_loader:
optimizer.zero_grad()
y = batch.y[:batch.batch_size]
y_hat = model(batch.x, batch.edge_index.to(device))[:batch.batch_size]
loss = F.cross_entropy(y_hat, y)
loss.backward()
optimizer.step()
total_loss += float(loss) * batch.batch_size
total_correct += int((y_hat.argmax(dim=-1) == y).sum())
total_examples += batch.batch_size
pbar.update(batch.batch_size)
pbar.close()
return total_loss / total_examples, total_correct / total_examples
@torch.no_grad()
def test():
model.eval()
y_hat = model.inference(data.x, subgraph_loader, device=device).argmax(dim=-1)
y = data.y.to(y_hat.device)
accs = []
for mask in [data.train_mask, data.val_mask, data.test_mask]:
accs.append(int((y_hat[mask] == y[mask]).sum()) / int(mask.sum()))
return accs
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# don't load data to GPU, as it will be loaded to GPU during sampling.
model, data, config = get_model_data_config(dataset_name, load_pretrained=False, device='cpu',
full_data=True)
model = model.to(device)
pretrained_file = f"{config['root_path']}/pretrained/{dataset_name}_pretrained.pt"
if os.path.exists(pretrained_file):
user_input = input('A pretrained file exist. Do you want to retrain? (y/n):')
if user_input.lower() != 'y':
print("Skipping training!")
sys.exit(0)
# Already send node features/labels to GPU for faster access during sampling:
data = data.to(device, 'x', 'y')
neig_args = config['nei_sampler_args']
kwargs = {'batch_size': neig_args['batch_size'], 'num_workers': 6, 'persistent_workers': True}
train_loader = NeighborLoader(data, input_nodes=data.train_mask,
num_neighbors=neig_args['sizes'], shuffle=True, **kwargs)
subgraph_loader = NeighborLoader(copy.copy(data), input_nodes=None,
num_neighbors=[-1], shuffle=False, **kwargs)
# No need to maintain these features during evaluation:
del subgraph_loader.data.x, subgraph_loader.data.y
# Add global node index information.
subgraph_loader.data.num_nodes = data.num_nodes
subgraph_loader.data.n_id = torch.arange(data.num_nodes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(1, 11):
loss, acc = train(epoch)
print(f'Epoch {epoch:02d}, Loss: {loss:.4f}, Approx. Train: {acc:.4f}')
train_acc, val_acc, test_acc = test()
print(f'Epoch: {epoch:02d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, '
f'Test: {test_acc:.4f}')
if test_acc > best_test:
best_test = test_acc
torch.save(model.state_dict(), pretrained_file)
torch.save(model.state_dict(), pretrained_file)
print(f"Model saved to {pretrained_file}.")
# Sample explain data and save. This is used for benchmarking explanation methods.
# This makes sure that the explain data is the same for all methods.
num_test_nodes = 100
explain_loader = NeighborLoader(data, input_nodes=data.test_mask.nonzero()[:num_test_nodes,0],
num_neighbors=[200, 50], batch_size=num_test_nodes,
num_workers=8, persistent_workers=True)
max_size = 0
max_ind = 0
avg_size = 0
batch = next(iter(explain_loader))
for i in range(batch.batch_size):
m = pruned_comp_graph(i, 2, batch.edge_index)[1].size(1)
if m > max_size:
max_size = m
max_ind = i
avg_size += m
del batch.x, batch.y # reduce saved file size in disk. Can be reloaded from the original data.
torch.save(batch, f"{config['root_path']}/pretrained/{dataset_name}_explain_data.pt")
print(f"Explain data saved to {config['root_path']}/pretrained.")
print("Maximum size: ", max_size, "max index: ", max_ind, "avg size: ",
avg_size / num_test_nodes)