-
Notifications
You must be signed in to change notification settings - Fork 1
/
train_GCN.py
154 lines (123 loc) · 5.94 KB
/
train_GCN.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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jul 16 20:01:29 2023
@author: chris
"""
import torch
import utility
import time
import numpy as np
import math
def full_batch_step(model, optimizer, criterion, device, x_train, y_train,
adj_train, train_mask, logging = False):
model.train()
optimizer.zero_grad()
out = model(x_train, adj_train)
if train_mask == None:
loss = criterion(out, y_train)
else:
loss = criterion(out[train_mask], y_train[train_mask])
loss.backward()
optimizer.step()
if logging:
acc,micro_f1,sens,spec = utility.metrics(out,y_train)
print(f"Train accuracy: {acc}, Train micro_f1: {micro_f1},Train Sens: {sens}, Train Spec: {spec}")
return loss
def evaluate(model, x, y, adj, mask):
with torch.no_grad():
model.eval()
out = model(x, adj)
acc,micro_f1,sens,spec = utility.metrics(out[mask],y[mask])
return acc, micro_f1, sens, spec
def train(model, device, x_train, y_train, adj_train, train_mask = None, x_val = None,
y_val = None, adj_val = None, val_mask = None, x_test = None,
y_test = None, adj_test = None, test_mask = None, multilabel = True,
lr = 0.0005, num_batch = 10, num_epoch = 100):
# passing model and training data to GPU
model = model.to(device)
x_train = x_train.to(device)
y_train = y_train.to(device)
adj_train = adj_train.to(device)
if train_mask != None:
train_mask = train_mask.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
if multilabel:
criterion = torch.nn.BCEWithLogitsLoss()
else:
criterion = torch.nn.CrossEntropyLoss()
max_val_acc = 0
max_val_sens = 0
max_val_spec = 0
max_val_f1 = 0
max_val_test_acc = 0
max_val_test_sens = 0
max_val_test_spec = 0
max_val_test_f1 = 0
time_arr = np.zeros((num_epoch,))
for epoch in range(num_epoch):
# single mini batch step
t = time.time()
loss = full_batch_step(model, optimizer, criterion, device,
x_train, y_train, adj_train, train_mask,
logging = False)
time_per_epoch = time.time() - t
time_arr[epoch] = time_per_epoch
if epoch == 0:
train_memory = torch.cuda.max_memory_allocated(device)*2**(-20)
# passing validation and test data to GPU (we do it after first forward pass to get)
# accurate pure training GPU memory usage
if x_val != None and y_val != None and adj_val != None and val_mask != None:
x_val = x_val.to(device)
y_val = y_val.to(device)
adj_val = adj_val.to(device)
val_mask = val_mask.to(device)
if x_test != None and y_test != None and adj_test != None and test_mask != None:
x_test = x_test.to(device)
y_test = y_test.to(device)
adj_test = adj_test.to(device)
test_mask = test_mask.to(device)
if epoch % 100 == 0:
print(f'Epoch: {epoch:03d}, Loss: {loss:.10f}, training time: {time_per_epoch:.5f}')
print(f"Peak GPU Memory Usage: {torch.cuda.max_memory_allocated(device)*2**(-20)} MB")
# evaluation
if x_val != None and y_val != None:
acc, micro_f1, sens, spec = evaluate(model, x_val, y_val, adj_val,
val_mask)
if epoch % 100 == 0:
print(f"Val accuracy: {acc}, Val micro_f1: {micro_f1}, Val Sens: {sens}, Val Spec: {spec}")
if acc > max_val_acc:
max_val_acc = acc
max_val_f1 = micro_f1
max_val_sens = sens
max_val_spec = spec
if (x_test != None and y_test != None):
acc, micro_f1, sens, spec = evaluate(model, x_test, y_test,
adj_test, test_mask)
max_val_test_acc = acc
max_val_test_f1 = micro_f1
max_val_test_sens = sens
max_val_test_spec = spec
print("===========================================Best Model Update:=======================================")
print(f"Val accuracy: {max_val_acc}, Val f1: {max_val_f1}, Val Sens: {max_val_sens}, Val Spec: {max_val_spec}")
print(f"Test accuracy: {max_val_test_acc}, Test f1: {max_val_test_f1}, Test Sens: {max_val_test_sens}, Test Spec: {max_val_test_spec}")
print("====================================================================================================")
print("Best Model:")
print(f"Val accuracy: {max_val_acc}, Val f1: {max_val_f1}, Val Sens: {max_val_sens}, Val Spec: {max_val_spec}")
print(f"Test accuracy: {max_val_test_acc}, Test f1: {max_val_test_f1}, Test Sens: {max_val_test_sens}, Test Spec: {max_val_test_spec}")
print(f"Average time per epoch: {time_arr[10:].mean()}") # don't include the first few epoch (slower due to Torch initialization)
print(f"Training GPU Memory Usage: {train_memory} MB")
print(f"Peak GPU Memory Usage: {torch.cuda.max_memory_allocated(device)*2**(-20)} MB")
# cleaning memory and stats
session_memory = torch.cuda.max_memory_allocated(device)*2**(-20)
train_time_avg = time_arr[10:].mean()
del x_val
del y_val
del x_test
del y_test
model = model.to('cpu')
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats(device)
return (max_val_acc, max_val_f1, max_val_sens, max_val_spec, max_val_test_acc,
max_val_test_f1, max_val_test_sens, max_val_test_spec, session_memory,
train_memory, train_time_avg)