def test_gnnmodel_build(self):
nodes = tf.one_hot([2, 4, 1, 3, 3], 16)
# neighbors -> 5 nodes, 2 neighbors
nlist = np.zeros((5, 2), dtype=np.int)
# make each atom neighbors with subsequent 2 atoms, mod 5
for i in range(5):
for k, j in enumerate(range(-1, 3, 2)):
nlist[i, k] = (i + j) % 5
nlist = tf.constant(nlist)
# 5 nodes, 2 neighbors, 4 feature
edges = tf.ones((5, 2))
# make inv degree (each one has two connections)
inv_degree = np.ones((5,)) / 2
inputs = [nodes, nlist, edges, inv_degree]
# make peak standards
ps = {}
for i in range(16):
ps[i] = ('F', 0, 1)
model = nmrgnn.build_GNNModel()
out_nodes = model(inputs)
# one peak per atom
assert out_nodes.shape == (nodes.shape[0],)
out_nodes = model(inputs)
assert out_nodes.shape == (nodes.shape[0],)
def test_model(self):
nodes = tf.one_hot([2, 4, 1, 3, 3], 16)
# neighbors -> 5 nodes, 2 neighbors
nlist = np.zeros((5, 2), dtype=np.int)
# make each atom neighbors with subsequent 2 atoms, mod 5
for i in range(5):
for k, j in enumerate(range(-1, 3, 2)):
nlist[i, k] = (i + j) % 5
nlist = tf.constant(nlist)
# 5 nodes, 2 neighbors, 4 feature
edges = tf.ones((5, 2))
# make inv degree (each one has two connections)
inv_degree = np.ones((5,)) / 2
inputs = [nodes, nlist, edges, inv_degree]
# make peak standards
ps = {}
for i in range(16):
ps[i] = ('F', 0, 1)
model = nmrgnn.build_GNNModel()
out_nodes = model(inputs)
try:
model.save('gnn_model_load_test')
del model
tf.keras.models.load_model(
'gnn_model_load_test', custom_objects=nmrgnn.custom_objects)
finally:
shutil.rmtree('gnn_model_load_test')
def test_edgeFCBlock_call(self):
model = nmrgnn.build_GNNModel()
hypers = model.hypers
edge_input = tf.ones((5, 2, 2))
edgeFCBlock = nmrgnn.EdgeFCBlock(hypers)
edge_output = edgeFCBlock(edge_input)
assert edge_output.shape[-1] == hypers.get('edge_feature_size')
assert edge_output.shape[:-1] == edge_input.shape[:-1]
def test_fcBlock_call(self):
model = nmrgnn.build_GNNModel()
hypers = model.hypers
nodes = tf.ones((5, hypers.get('atom_feature_size')))
fcBlock = nmrgnn.FCBlock(hypers)
new_nodes = fcBlock(nodes)
assert new_nodes.shape[-1] == hypers.get('atom_feature_size') // 2
assert new_nodes.shape[-1] == nodes.shape[-1] // 2
def train(tfrecords, name, epochs, embeddings, validation, checkpoint_path,
tensorboard, load, loss_balance):
'''Train the model'''
model = nmrgnn.build_GNNModel(loss_balance=loss_balance)
if load:
model.load_weights(checkpoint_path)
callbacks = []
# set-up learning rate scheduler
reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss',
factor=0.99,
patience=4,
min_lr=1e-4,
verbose=1)
callbacks.append(reduce_lr)
# tensorboard
if tensorboard is not None:
tensorboard_callback = tf.keras.callbacks.TensorBoard(
log_dir=tensorboard,
write_images=False,
write_graph=False,
histogram_freq=0,
profile_batch=0)
callbacks.append(tensorboard_callback)
# save model
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
filepath=checkpoint_path,
save_weights_only=True,
monitor='val_loss',
save_best_only=False)
callbacks.append(model_checkpoint_callback)
train_data, validation_data = load_data(tfrecords,
validation,
embeddings,
scale=False)
# explicitly call model to get shapes defined
for t in train_data:
x, y, m = t
model(x)
break
results = model.fit(train_data,
epochs=epochs,
callbacks=callbacks,
validation_data=validation_data,
validation_freq=1)
model.save(name)
pfile = name + '-history-0.pb'
i = 0
while os.path.exists(pfile):
i += 1
pfile = f'{name}-history-{i}.pb'
with open(pfile, 'wb') as f:
pickle.dump(results.history, file=f)
def test_mpBlock_call(self):
model = nmrgnn.build_GNNModel()
hypers = model.hypers
nodes = tf.one_hot([2, 0, 1, 3, 3], 16)
# neighbors -> 5 nodes, 2 neighbors
nlist = np.zeros((5, 2), dtype=np.int)
# make each atom neighbors with subsequent 2 atoms, mod 5
for i in range(5):
for k, j in enumerate(range(-1, 3, 2)):
nlist[i, k] = (i + j) % 5
nlist = tf.constant(nlist)
# 5 nodes, 2 neighbors, 4 feature
edges = tf.ones((5, 2, 2))
# make inv degree (each one has two connections)
inv_degree = np.ones((5,)) / 2
# shapes are specified inside the block
mp_block = nmrgnn.MPBlock(hypers)
out_nodes = mp_block([nodes, nlist, edges, inv_degree])
assert out_nodes.shape == nodes.shape