Force training with angular symmetry functions for clusters with different number of atoms

[lmj1029123]

Hi,

It seems that the trainer is not working properly when I did force training with angular symmetry functions for clusters with different number of atoms. (ElementalEnergy)

When I did the training for clusters with same number of atoms, it is OK. By looking at the log, it seems that the losses are going down. However, if I use it for clusters with different number of atoms, the losses become nan. This applies to both "Behler" and "Weighted" mode. Changing the learning rate of the Adam optimizer does not help. I don't know exactly why this is the case. My guess is that it could be due to the padding you do to the "representation" when you have clusters of different size. Could you have a look at that? Thanks!

from schnetpack.data import AtomsData
import torch
import torch.nn.functional as F
from torch.optim import Adam
from torch.optim import LBFGS
import pickle 
import schnetpack as spk
import schnetpack.atomistic as atm
import schnetpack.representation as rep
from schnetpack.datasets import *
import schnetpack.evaluation as eva
from schnetpack.metrics import MeanSquaredError

Name = 'debug'  # (ref:Name)

if not os.path.exists(Name):
    os.makedirs(Name)

data = AtomsData('./db/master-jb.db', required_properties=['energy','forces'],collect_triples=True)

#40 training data, 5 val data, rest are for test
train, val, test = data.create_splits(30, 30)
train_loader = spk.data.AtomsLoader(train, batch_size=3, num_workers=0)
val_loader = spk.data.AtomsLoader(val)
test_loader = spk.data.AtomsLoader(test)
loader = [train_loader, val_loader, test_loader]
pickle.dump(loader, open('./'+Name+'/loader.sav','wb'))

reps = rep.BehlerSFBlock(n_radial=2, n_angular=2, elements=frozenset([46,79]), cutoff_radius=6.0, mode = 'weighted')
#print(reps.n_symfuncs)
output = ElementalEnergy(n_in=reps.n_symfuncs,n_layers=3,n_hidden=10, elements=frozenset([46,79]),return_force=True,create_graph=True)
model = atm.AtomisticModel(reps, output)

trainable_params = filter(lambda p: p.requires_grad, model.parameters())

metric_E = [MeanSquaredError(target = 'energy',model_output='y'),MeanSquaredError(target = 'forces',model_output='dydx')]
hook_E = spk.train.CSVHook("./"+Name, metric_E)

opt = Adam(trainable_params, lr=1e-4)

loss = lambda b, p: F.mse_loss(p["y"], b['energy'])+F.mse_loss(p["dydx"], b['forces'])
trainer = spk.train.Trainer(Name+"/", model, loss,
                      opt, train_loader, val_loader,hooks = [hook_E])

# start training
trainer.train(torch.device("cpu"))

Best, Mingjie

[mgastegger]

I found the problem and it should be fixed now.

On a side note, I would also recommend adding a special standardization layer for symmetry functions and one for the energies. For the SF, this would look something like this:

reps = rep.BehlerSFBlock(...)
reps = rep.StandardizeSF(reps, train_loader)

The scaling for the energy in the output layer can be done via:

mean, stddev = train_loader.get_statistics('energy', True)
output = atm.ElementalEnergy(...
                             create_graph=True,
                             mean=mean,
                             stddev=stddev)

Hope this could fix your issue.