Optimizing the training

# install konda
!pip -q install konda

import konda
konda.install()

# accept Anaconda Terms of Service
!conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
!conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r

# create the environment
!konda create -q -n pinn_env python=3.10 -y
!konda activate pinn_env

# install PiNN & download the QM9 dataset
!konda run "pip -q install 'tensorflow[and-cuda]==2.15.1' 'numpy<2'"
!konda run "pip -q 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

%%writefile imports.py

import sys
import os, warnings
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import tensorflow as tf

import timeit
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

os.environ['CUDA_VISIBLE_DEVICES'] = ''
index_warning = 'Converting sparse IndexedSlices'
warnings.filterwarnings('ignore', index_warning)

%%writefile -a imports.py

# For the purpose of testing, we use only 1000 samples from QM9
filelist = glob('/tmp/dsgdb9nsd/*.xyz')[:1000]
dataset = lambda: load_qm9(filelist)

!rm -f run.py
!cat imports.py >> run.py

%%writefile -a run.py

ds = dataset().repeat().apply(sparse_batch(100))
tensors = ds.as_numpy_iterator()
for i in range(10):
    next(tensors) # "Warm up" the graph

print(f"{timeit.timeit(lambda: next(tensors), number=100) / 100 * 1e3:.1f} ms per iteration, average of 100")
sys.exit()

!konda run "python run.py"

!rm -f run_cached.py
!cat imports.py >> run_cached.py

%%writefile -a run_cached.py

ds = dataset().cache().repeat().apply(sparse_batch(100))
tensors = ds.as_numpy_iterator()
for i in range(10):
    next(tensors) # "Warm up" the graph
print(f"{timeit.timeit(lambda: next(tensors), number=100) / 100 * 1e3:.1f} μs per iteration, average of 100")
sys.exit()

!konda run "python run_cached.py"

!rm -f run.py
!cat imports.py >> run.py

%%writefile -a run.py

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

print(f"{timeit.timeit(lambda: next(tensors), number=100) / 100 * 1e3:.1f} ms per per iteration, average of 100")
sys.exit()

!konda run "python run.py"

!rm -f run_preprocessed.py
!cat imports.py >> run_preprocessed.py

%%writefile -a run_preprocessed.py

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
print(f"{timeit.timeit(lambda: next(tensors), number=100) / 100 * 1e3:.1f} μs per iteration, average of 100")
sys.exit()

!konda run "python run_preprocessed.py"

!rm -f run_preprocessed_cached.py
!cat imports.py >> run_preprocessed_cached.py

%%writefile -a run_preprocessed_cached.py

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

print(f"{timeit.timeit(lambda: next(tensors), number=100) / 100 * 1e3:.1f} μs per iteration")
sys.exit()

!konda run "python run_preprocessed_cached.py"

!rm -f atomic_dress.py
!cat imports.py >> atomic_dress.py

%%writefile -a atomic_dress.py

import numpy as np
import json

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])

with open("dress.json", "w") as f:
  json.dump(dress, f)
np.save("error.npy", error)

#adding a dress to the whole QM9 dataset. running this cell might take some time.
!konda run "python atomic_dress.py"

import json
import numpy as np
import matplotlib.pyplot as plt

with open("dress.json", "r") as f:
  dress = json.load(f)
error = np.load("error.npy")

plt.hist(error,50)
dress

%%writefile training.py

from glob import glob
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import tensorflow as tf

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

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])

%%writefile -a training.py

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)

%%writefile -a training.py

tf.estimator.train_and_evaluate(model, train_spec, eval_spec)

!rm -rf /tmp/PiNet_QM9_pipeline
!konda run "python training.py"

# suppose you have two cards
distribution = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
config = tf.estimator.RunConfig(train_distribute=distribution)