Training a potential from scratch

[ligerzero-ai]

Nvm -

The correct way to do this is (?):

import pandas as pd
import warnings

import numpy as np
from pymatgen.core import Structure



# To suppress warnings for clearer output
warnings.simplefilter("ignore")

df = pd.read_pickle("train_dataset.pkl")
#df.structure.iloc[0]
df["forces"] = [row.forces.tolist() for _, row in df.iterrows()]

df["structure_reconst"] = [Structure.from_str(row.structure, fmt="json") for _, row in df.iterrows()]
structures = df.structure_reconst.tolist()
stresses = [np.zeros((3, 3)).tolist() for s in structures]

from chgnet.data.dataset import StructureData, get_train_val_test_loader
from chgnet.trainer import Trainer

dataset = StructureData(
    structures=structures,
    energies=df.energy.tolist(),
    forces=df.forces.tolist(),
#    stresses=list_of_stresses,
#    magmoms=list_of_magmoms,
)
train_loader, val_loader, test_loader = get_train_val_test_loader(
    dataset, batch_size=32, train_ratio=0.9, val_ratio=0.05
)

# Optionally fix the weights of some layers
#for layer in [
#    chgnet.atom_embedding,
#    chgnet.bond_embedding,
#    chgnet.angle_embedding,
#    chgnet.bond_basis_expansion,
#    chgnet.angle_basis_expansion,
#    chgnet.atom_conv_layers[:-1],
#    chgnet.bond_conv_layers,
#    chgnet.angle_layers,
#]:
#    for param in layer.parameters():
#        param.requires_grad = False
from chgnet.model import CHGNet

# Load pretrained CHGNet
chgnet = CHGNet()

# Define Trainer
trainer = Trainer(
    model=chgnet,
    targets="ef",
    optimizer="Adam",
    scheduler="CosLR",
    criterion="MSE",
    epochs=5000,
    learning_rate=1e-2,
    use_device="cuda",
    print_freq=6,
)
trainer.train(train_loader, val_loader, test_loader)

[rydeveraumn]

@ligerzero-ai did you successfully train from scratch on the full dataset? How long did it take and how many epoch was the optimal number?

[ligerzero-ai]

Hey @rydeveraumn I didn't manage to train it, it was put on the backburner - I did have some success with MACE/allegro potentials, however. My dataset (magnetic Fe-X) structures wasn't amenable to the as-set loss function and I didn't bother after that. MACE supports multi-gpu training out of the box now, so I went with that instead. MACE github

[rydeveraumn]

Hey @ligerzero-ai thank you so much for the response and the link! I will take a look. I have been referred to MACE 2 times now 🙂

[rydeveraumn]

@bowen-bd I still have some questions related to training from scratch, I am connected with Chris Bartel and he said you might be able to help. I downloaded the full MPtrj_2022.9_full.json and I see the notes in the figshare:

3. There're missing MAGMOMs labels in the MPtrj, which we put None as labels. These do not mean the MAGMOM label is 0. CHGNet is trained on absolute value of DFT magmom, which is the absolute value of the labels contained in MPtrj, the unit conversion is automatic if you use the dataset we provide, see: https://github.com/CederGroupHub/chgnet/blob/main/chgnet/data/dataset.py

My question is if the model was trained on the full dataset and used a mask loss for these missing magmoms or if the model was only trained on structures with magmom data.

[bowen-bd]

the structures without magmom are also included in the training, just with the magmom loss masked out.