Fine-tuning CHGNet with AIMD data

Science Talk
, , ,
1

I was wondering if anyone has tried to fine-tune the recent CHGNet graph neural network potential. I’m having some basic use questions for how to fine-tune using some AIMD data. The README file has the following:

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

dataset = StructureData(
    structures=list_of_structures,
    energies=list_of_energies,
    forces=list_of_forces,
    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
)
trainer = Trainer(
    model=chgnet,
    targets='efsm',
    optimizer='Adam',
    criterion='MSE',
    learning_rate=1e-2,
    epochs=50,
    use_device='cuda'
)

trainer.train(train_loader, val_loader, test_loader)

The first question I have is what are the formats/data-types of the structures, energies, etc. in the lists passed to StructureData? Are they JSON? The second question is how do you save the weights and load them for later use? The third question is with regard to the note:

  1. The pretrained dataset of CHGNet comes from GGA+U DFT with MaterialsProject2020Compatibility. The parameter for VASP is described in MPRelaxSet. If you’re fine-tuninig with MPRelaxSet, it is recommended that apply the MP2020 compatibility to your energy labels so that they’re consistent with the pretrained dataset.

I’m not sure what the implications are on GGA AIMD simulations data using Nose-Hoover thermostat.

If anyone has put together a Jupyter notebook that demonstrates fine-tuning using some AIMD data that would be extremely useful.

2

Your 1st question might have been addressed by StructureData types and docs by janosh · Pull Request #17 · CederGroupHub/chgnet · GitHub

For your 2nd question, I’d use torch.save() and torch.load().

I believe that’s in the works and should be coming soon.

3

Hi, Stefan,

please see our recent update in the GitHub, which includes a example notebook of fine tuning.

If you have further concerns, please initiate an issue in the GitHub.

1 Like
4

@janosh thanks for the comments, helped me get me navigate.

@Bowen_Deng will definitely look through the new updates. Thanks!