The inconsistent property key for Multihead

[MengnanCui]

Describe the bug
Related to the issue
When I do Multihead transfer learning from a personal pre-trained model. I found that the energy_key and forces_key should be the key for pre-training properties(dftb_energy, dftb_forces in this case) rather than dft_energy and dft_forces from args.energy_key and args.forces_key. Otherwise it can not read corrected properties for pt_head.
line 288 .. 323

collections, atomic_energies_dict = get_dataset_from_xyz(
                work_dir=args.work_dir,
                train_path=args.pt_train_file,
                valid_path=args.pt_valid_file,
                valid_fraction=args.valid_fraction,
                config_type_weights=None,
                test_path=None,
                seed=args.seed,
                energy_key='dftb_energy',#args.energy_key,
                forces_key='dftb_forces',#args.forces_key,
                stress_key=args.stress_key,
                virials_key=args.virials_key,
                dipole_key=args.dipole_key,
                charges_key=args.charges_key,
                head_name="pt_head",
                keep_isolated_atoms=args.keep_isolated_atoms,
            )...

To Reproduce

1. when use

energy_key=args.energy_key,
forces_key=args.forces_key,

2024-11-13 13:12:30.501 INFO: ==================Using multiheads finetuning mode==================
2024-11-13 13:12:30.501 INFO: Using foundation model for multiheads finetuning with ../train_pt.xyz
2024-11-13 13:12:32.778 INFO: Training set [7642 configs, 0 energy, 380589 forces] loaded from '../train_pt.xyz'
2024-11-13 13:12:33.279 INFO: Validation set [1000 configs, 1000 energy, 46593 forces] loaded from '../valid_pt.xyz'
2024-11-13 13:12:33.279 INFO: Total number of configurations: train=7642, valid=1000

2. after changing the key

energy_key='dftb_energy',#args.energy_key,
forces_key='dftb_forces',#args.forces_key,

2024-11-13 13:12:30.501 INFO: ==================Using multiheads finetuning mode==================
2024-11-13 13:12:30.501 INFO: Using foundation model for multiheads finetuning with ../train_pt.xyz
2024-11-13 13:12:32.778 INFO: Training set [7642 configs, 7642 energy, 380589 forces] loaded from '../train_pt.xyz'
2024-11-13 13:12:33.279 INFO: Validation set [1000 configs, 1000 energy, 46593 forces] loaded from '../valid_pt.xyz'
2024-11-13 13:12:33.279 INFO: Total number of configurations: train=7642, valid=1000

[ilyes319]

Hey @MengnanCui, right now the two set of keys should match exactly.

[MengnanCui]

Hi, Ilyes. I reinstall the main branch. but the outputs seems to be the same as shown above, the energy_key and forces_key for pre-trained model are still missing.

I compared it with the multihead fine-tunning based on the foundation model,

After the line 286 of run_train.py

            head_config_pt = HeadConfig(
                head_name="pt_head",
                E0s="foundation",
                statistics_file=args.statistics_file,
                compute_avg_num_neighbors=False,
                avg_num_neighbors=model_foundation.interactions[0].avg_num_neighbors,
            )

the energy_key and forces_key are defined for the foundation model inside of the head_config_pt, as shown in the following (None for both property keys):

HeadConfig(head_name='pt_head', train_file=None, valid_file=None, test_file=None, test_dir=None, E0s='foundation', statistics_file=None, valid_fraction=None, config_type_weights=None, energy_key=None, forces_key=None, stress_key=None, virials_key=None, dipole_key=None, charges_key=None, keep_isolated_atoms=None, atomic_numbers=None, mean=None, std=None, avg_num_neighbors=61.964672446250916, compute_avg_num_neighbors=False, collections=None, train_loader=None, z_table=None, atomic_energies_dict=None)

where these arguments are defined in multihead_tools.py manually if I understood correctly, so I guess the property keys for pre-trained model have to be defined manually while doing multihead fine-tunning?

Also, I did not find the property keys are saved inside of pre-trained models, maybe this is the reason why it cant do?

Thank you and kind regards from Mengnan