Optimizing the training 
This notebooks covers more details on tweaking and optimizing the training process.
# Install PiNN & download QM9 dataset
!pip install tensorflow==2.9
!pip install git+https://github.com/Teoroo-CMC/PiNN
!mkdir -p /tmp/dsgdb9nsd && curl -sSL https://ndownloader.figshare.com/files/3195389 | tar xj -C /tmp/dsgdb9nsd
import os, warnings
import tensorflow as tf
import matplotlib.pyplot as plt
from glob import glob
from pinn.io import load_qm9, sparse_batch
from pinn.networks.pinet import PiNet
from pinn.utils import get_atomic_dress
from pinn import get_model, get_network
os.environ['CUDA_VISIBLE_DEVICES'] = ''
index_warning = 'Converting sparse IndexedSlices'
warnings.filterwarnings('ignore', index_warning)
Optimizing the pipeline
Caching
Caching stores the decoded dataset in the memory.
# For the purpose of testing, we use only 1000 samples from QM9
filelist = glob('/tmp/dsgdb9nsd/*.xyz')[:1000]
dataset = lambda: load_qm9(filelist)
ds = dataset().repeat().apply(sparse_batch(100))
tensors = ds.as_numpy_iterator()
for i in range(10):
next(tensors) # "Warm up" the graph
%timeit next(tensors)
This speed indicates the IO limit of our current setting.
Now let's cache the dataset to the memory.
ds = dataset().cache().repeat().apply(sparse_batch(100))
tensors = ds.as_numpy_iterator()
for i in range(10):
next(tensors) # "Warm up" the graph
%timeit next(tensors)
Preprocessing
You might also see a notable difference in the performance with and without preprocessing. This is especially helpful when you are training with GPUs.
pinet = PiNet()
ds = dataset().cache().repeat().apply(sparse_batch(100))
tensors = ds.as_numpy_iterator()
for i in range(10):
pinet(next(tensors)) # "Warm up" the graph
%timeit pinet(next(tensors))
pinet = PiNet()
ds = dataset().cache().repeat().apply(sparse_batch(100)).map(pinet.preprocess)
tensors = ds.as_numpy_iterator()
for i in range(10):
next(tensors) # "Warm up" the graph
%timeit next(tensors)
You can even cache the preprocessed data.
pinet = PiNet()
ds = dataset().apply(sparse_batch(100)).map(pinet.preprocess).cache().repeat()
tensors = ds.as_numpy_iterator()
for i in range(10):
next(tensors) # "Warm up" the graph
%timeit next(tensors)
Atomic dress
Scaling and aligning the labels can enhance the performance of the models, and avoid numerical instability. For datasets like QM9, we can assign an atomic energy to each atom according to their elements to approximate the total energy. This can be done by a simple linear regression. We provide a simple tool to generate such "atomic dresses".
filelist = glob('/tmp/dsgdb9nsd/*.xyz')
dataset = lambda: load_qm9(filelist, splits={'train':8, 'test':2})
dress, error = get_atomic_dress(dataset()['train'],[1,6,7,8,9])
Applying the atomic dress converts the QM9 energies to a "normal" distribution. It also gives us some ideas about the relative distribution of energies, and how much our neural network improves from the naive guess of the atomic dress.
After applying the atomic dress, it turns out that the distribution of our training set is only about 0.05 Hartree, or 30 kcal/mol.
plt.hist(error,50)
dress
Training with the optimized pipeline
!rm -rf /tmp/PiNet_QM9_pipeline
params = {'model_dir': '/tmp/PiNet_QM9_pipeline',
'network': {
'name': 'PiNet',
'params': {
'atom_types':[1, 6, 7, 8, 9],
},
},
'model': {
'name': 'potential_model',
'params': {
'learning_rate': 1e-3, # Relatively large learning rate
'e_scale': 627.5, # Here we scale the model to kcal/mol
'e_dress': dress
}
}
}
# The logging behavior of estimator can be controlled here
config = tf.estimator.RunConfig(log_step_count_steps=500)
# Preprocessing the datasets
model = get_model(params, config=config)
# If you are pre-processing the dataset in the training script,
# the preprocessing layer will occupy the namespace of the network
# resulting unexpected names in the ckpts and errors durning prediction
# To avoid this, wrap your preprocessing function with a name_scope.
# This will not be a problem if you save a preprocessed dataset
def pre_fn(tensors):
with tf.name_scope("PRE") as scope:
network = get_network(model.params['network'])
tensors = network.preprocess(tensors)
return tensors
train = lambda: dataset()['train'].apply(sparse_batch(100)).map(pre_fn).cache().repeat().shuffle(100)
test = lambda: dataset()['test'].apply(sparse_batch(100))
# Running specs
train_spec = tf.estimator.TrainSpec(input_fn=train, max_steps=1e4)
eval_spec = tf.estimator.EvalSpec(input_fn=test, steps=100)
tf.estimator.train_and_evaluate(model, train_spec, eval_spec)
Monitoring
It's recommended to monitor the training with Tensorboard instead of the stdout here.
Try tensorboard --logdir /tmp
Parallelization with tf.Estimator
The estimator api makes it extremely easy to train on multiple GPUs.
# suppose you have two cards
distribution = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
config = tf.estimator.RunConfig(train_distribute=distribution)
Conclusions
Congratulations! You can now train atomic neural networks with state-of-the-art accuracy and speed.
But there's more. With PiNN, the components of ANNs are modulized. Read the following notebooks to see how you can build your own ANN.