-
Notifications
You must be signed in to change notification settings - Fork 0
/
graph_neural_network.py
128 lines (103 loc) · 4.19 KB
/
graph_neural_network.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
import numpy as np
import tensorflow as tf
from sklearn.model_selection import train_test_split
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.losses import MeanSquaredError
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from spektral.datasets import qm9
from spektral.layers import EdgeConditionedConv, ops, GlobalSumPool
from spektral.utils import batch_iterator, numpy_to_disjoint
from spektral.utils import label_to_one_hot
learning_rate = 1e-3 # Learning rate
epochs = 100 # Number of training epochs
batch_size = 256 # Batch size
tf.keras.backend.set_floatx('float16')
# Also check eager vs non eager execution
A, X, E, y = qm9.load_data(return_type='numpy',
nf_keys='atomic_num',
ef_keys='type',
self_loops=True,
auto_pad=False,
amount=133000) # Set to None to train on whole dataset
y = y[['cv']].values # Heat capacity at 298.15K
# Preprocessing
X_uniq = np.unique([v for x in X for v in np.unique(x)])
E_uniq = np.unique([v for e in E for v in np.unique(e)])
X_uniq = X_uniq[X_uniq != 0]
E_uniq = E_uniq[E_uniq != 0]
X = [label_to_one_hot(x, labels=X_uniq) for x in X]
E = [label_to_one_hot(e, labels=E_uniq) for e in E]
# Parameters
F = X[0].shape[-1] # Dimension of node features
S = E[0].shape[-1] # Dimension of edge features
n_out = y.shape[-1] # Dimension of the target
A_train, A_test, \
X_train, X_test, \
E_train, E_test, \
y_train, y_test = train_test_split(A, X, E, y, test_size=0.1, random_state=0)
X_in = Input(shape=(F,), name='X_in')
A_in = Input(shape=(None,), sparse=True, name='A_in')
E_in = Input(shape=(S,), name='E_in')
I_in = Input(shape=(), name='segment_ids_in', dtype=tf.int32)
X_1 = EdgeConditionedConv(32, activation='relu')([X_in, A_in, E_in])
X_2 = EdgeConditionedConv(32, activation='relu')([X_1, A_in, E_in])
X_3 = GlobalSumPool()([X_2, I_in])
output = Dense(n_out)(X_3)
# Build model
model = Model(inputs=[X_in, A_in, E_in, I_in], outputs=output)
opt = Adam(lr=learning_rate)
loss_fn = MeanSquaredError()
@tf.function(
input_signature=(tf.TensorSpec((None, F), dtype=tf.float64),
tf.SparseTensorSpec((None, None), dtype=tf.float64),
tf.TensorSpec((None, S), dtype=tf.float64),
tf.TensorSpec((None,), dtype=tf.int32),
tf.TensorSpec((None, n_out), dtype=tf.float64)),
experimental_relax_shapes=True)
def train_step(X_, A_, E_, I_, y_):
with tf.GradientTape() as tape:
predictions = model([X_, A_, E_, I_], training=True)
loss = loss_fn(y_, predictions)
loss += sum(model.losses)
gradients = tape.gradient(loss, model.trainable_variables)
opt.apply_gradients(zip(gradients, model.trainable_variables))
return loss
import time
tic = time.perf_counter()
current_batch = 0
model_loss = 0
batches_in_epoch = np.ceil(len(A_train) / batch_size)
print('Fitting model')
batches_train = batch_iterator([X_train, A_train, E_train, y_train],
batch_size=batch_size, epochs=epochs)
for b in batches_train:
X_, A_, E_, I_ = numpy_to_disjoint(*b[:-1])
A_ = ops.sp_matrix_to_sp_tensor(A_)
y_ = b[-1]
outs = train_step(X_, A_, E_, I_, y_)
model_loss += outs.numpy()
current_batch += 1
if current_batch == batches_in_epoch:
print('Loss: {}'.format(model_loss / batches_in_epoch))
model_loss = 0
current_batch = 0
toc = time.perf_counter()
duration = toc - tic
print("Training time took {0}".format(duration))
tic = time.perf_counter()
print('Testing model')
model_loss = 0
batches_in_epoch = np.ceil(len(A_test) / batch_size)
batches_test = batch_iterator([X_test, A_test, E_test, y_test], batch_size=batch_size)
for b in batches_test:
X_, A_, E_, I_ = numpy_to_disjoint(*b[:-1])
A_ = ops.sp_matrix_to_sp_tensor(A_)
y_ = b[3]
predictions = model([X_, A_, E_, I_], training=False)
model_loss += loss_fn(y_, predictions)
model_loss /= batches_in_epoch
print('Done. Test loss: {}'.format(model_loss))
toc = time.perf_counter()
duration = toc - tic
print("Testing time took {0}".format(duration))