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)