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.