CHGNet training crash while importing json data that is converted from VASP AIMD.

[turbosonics]

I hope to conduct a test to train my own VASP AIMD data with or without pre-trained data. As a first step, I performed conversion of VASP AIMD to json file using pymatgen, same way as the "fine_tuning" example:

import numpy as np
from pymatgen.core import Structure
from chgnet.utils import parse_vasp_dir
dataset_dict = parse_vasp_dir(file_root="${INPUT_DIR}/${KEYWORD}")
print(list(dataset_dict))

# Structure to json
from chgnet.utils import write_json
dict_to_json = [struct.as_dict() for struct in dataset_dict["structure"]]
write_json(dict_to_json, "${OUTPUT_DIR}/${SIMULATION}/${KEYWORD}.json")

Then I used result json file to training, following the fine tuning example script:

import os
#import pickle as pk
from chgnet.model import CHGNet
from chgnet.data.dataset import StructureData, collate_graphs, StructureJsonData, get_train_val_test_loader
from chgnet.graph.converter import CrystalGraphConverter
from chgnet.trainer import Trainer
from pymatgen.core.structure import Structure
from torch.utils.data import DataLoader, Dataset

input_geo = '${GEOMETRY_DIR}/${GEOMETRY_TRAIN_NAME}'
graph_converter = CrystalGraphConverter(atom_graph_cutoff=6, bond_graph_cutoff=4)
dataset = StructureJsonData(
    input_geo,
    graph_converter,
    targets="efsm"
)
train_loader, val_loader, test_loader = get_train_val_test_loader(
    dataset, batch_size=10, train_ratio=0.9, val_ratio=0.05
)
...

But now I see the error:

...
Importing: /home/Geometry_storage/CHGNet_tests/01a_NPT-300K-magmom_pymatgen_VASPtoJSON/testNPT.json
Traceback (most recent call last):
  File "/mnt/work/CHGNet_training/CHGNet_test/01a1_test_json/test_fromscratch.py", line 15, in <module>
    dataset = StructureJsonData(
  File "/home/venv_chgnet_gpunode/lib/python3.9/site-packages/chgnet/data/dataset.py", line 516, in __init__
    self.data.update(utils.read_json(data))
ValueError: dictionary update sequence element #0 has length 6; 2 is required

Just in case, I also tried json file that has been converted from the same VASP AIMD data using ASE, but this time, CHGNet printed different error. (One thing to note is that the format of json file converted by ASE is different from json file converted by pymatgen.)

...
Importing: /home/Geometry_storage/CHGNet_tests/01b_NPT-300K-magmom_ASE_VASPtoJSON/testNPT_ASE.json
Traceback (most recent call last):
  File "/mnt/work/CHGNet_training/CHGNet_test/01a2_test_ASEjson/test_fromscratch.py", line 15, in <module>
    dataset = StructureJsonData(
  File "/home/venv_chgnet_gpunode/lib/python3.9/site-packages/chgnet/data/dataset.py", line 522, in __init__
    self.keys = [
  File "/home/venv_chgnet_gpunode/lib/python3.9/site-packages/chgnet/data/dataset.py", line 523, in <listcomp>
    (mp_id, graph_id) for mp_id, dct in self.data.items() for graph_id in dct
TypeError: 'int' object is not iterable

It seems that converted json file is the problem but I don't know what is wrong. Am I missing something during VASP AIMD data conversion or set up the training?

==========================

Just in case to add, I even tried to download the CHGNet source code using git clone and performed pip install, yet the same thing happens. So far, when I compile CHGNet, I used virtual environment with python 3.9. Should I need to try manualy compile specific version of pytorch? Our cluster have several version of Cuda but I loaded Cuda 11.8 when I compile CHGNet.

[BowenD-UCB]

Can you please check whether the structures and labels in the parsed json file is valid?
If they seem OK, can you try passing them into the most basic StructureData dataset?

[turbosonics]

Hello, thanks for your reply.

I think the the json file converted by parse_vasp_dir and write_json from VASP output files are not properly parsed. It looks different from typical json example. Also, as I mentioned, json file converted from the parse_vasp_dir and write_json looks totally different from json file converted using ASE.

StructureData didn't work. It showed following error:

...
Traceback (most recent call last):
  File "/mnt/work/CHGNet_training/CHGNet_test/01a2_test_json/test_fromscratch2.py", line 15, in <module>
    dataset = StructureData(
TypeError: __init__() got an unexpected keyword argument 'targets'

[BowenD-UCB]

This has been further addressed in 6408de5
Please checkout the updated fine-tuning.ipynb

[turbosonics]

Hello,

Thanks for your reaction and answer for question I raised. But I found some bugs(?) from your fine-tuning.ipynb

From 1. Prepare Training Data section,

stresses = None if result_dict['stress'] in [None, []] else result_dict['stress']
magmoms = None if result_dict['magmom'] in [None, []] else result_dict['magmom']

I think this should be

stresses = None if dataset_dict['stress'] in [None, []] else dataset_dict['stress']
magmoms = None if dataset_dict['magmom'] in [None, []] else dataset_dict['magmom']

right?

Also, from 2. Define Dataset

dataset = StructureData(
    structures=structures,
    energies=energies_per_atom,
    forces=forces,
    stresses=stresses,  # can be None
    magmoms=magmoms,  # can be None

This should be

dataset = StructureData(
    structures=structures,
    energies=energies,
    forces=forces,
    stresses=stresses,  # can be None
    magmoms=magmoms,  # can be None

Let me know if these corrections are OK. Thank you!