RuntimeError: shape '[128, 3, 128]' is invalid for input of size 98304

[YaoqiPang]

Hi,
Recently, I have been trying to use finetune function of MACE, I trianed a model from scratch, and then loaded this model as foundation_model to continue training with a expanded dataset. For finetuning, I only changed max_num_epochs and add a extra line: --foundation_model="./old/model-3_swa.model" . Besides that, nothing changes. But it gave me this error: warnings.warn(
Traceback (most recent call last):
File "/users/ypang12/.conda/envs/mace/bin/mace_run_train", line 8, in
sys.exit(main())
File "/users/ypang12/.conda/envs/mace/lib/python3.10/site-packages/mace/cli/run_train.py", line 66, in main
run(args)
File "/users/ypang12/.conda/envs/mace/lib/python3.10/site-packages/mace/cli/run_train.py", line 572, in run
model, output_args = configure_model(args, train_loader, atomic_energies, model_foundation, heads, z_table)
File "/users/ypang12/.conda/envs/mace/lib/python3.10/site-packages/mace/tools/model_script_utils.py", line 133, in configure_model
model = load_foundations_elements(
File "/users/ypang12/.conda/envs/mace/lib/python3.10/site-packages/mace/tools/finetuning_utils.py", line 82, in load_foundations_elements
.skip_tp.weight.reshape(
RuntimeError: shape '[128, 3, 128]' is invalid for input of size 98304.

Here is my initial train script:
mace_run_train
--name=model-3
--train_file="train.xyz"
--batch_size=64
--valid_fraction=0.1
--eval_interval=1
--patience=2000
--model="MACE"
--E0s="average"
--compute_avg_num_neighbors=True
--hidden_irreps="128x0e + 128x1o"
--num_interactions=3
--num_cutoff_basis=5
--num_radial_basis=8
--correlation=3
--r_max=12
--max_ell=3
--max_num_epochs=2000
--loss="weighted"
--amsgrad
--optimizer="adam"
--lr=0.005
--lr_factor=0.5
--scheduler="ReduceLROnPlateau"
--scheduler_patience=30
--scaling="rms_forces_scaling"
--error_table="PerAtomRMSE"
--default_dtype="float64"
--restart_latest
--device=cuda
--seed=1615767616
--energy_key="energy"
--forces_key="forces"
--energy_weight=5.0
--forces_weight=10.0
--ema
--ema_decay=0.95
--swa
--swa_lr=0.001
--swa_energy_weight=10.0
--swa_forces_weight=5.0
Here is my finetun script:
mace_run_train
--name=model-3
--train_file="train.xyz"
--batch_size=64
--valid_fraction=0.1
--eval_interval=1
--patience=2000
--model="MACE"
--foundation_model="./old/model-3_swa.model"
--E0s="average"
--compute_avg_num_neighbors=True
--hidden_irreps="128x0e + 128x1o"
--num_interactions=3
--num_cutoff_basis=5
--num_radial_basis=8
--correlation=3
--r_max=12
--max_ell=3
--max_num_epochs=1000
--loss="weighted"
--amsgrad
--optimizer="adam"
--lr=0.005
--lr_factor=0.5
--scheduler="ReduceLROnPlateau"
--scheduler_patience=30
--scaling="rms_forces_scaling"
--error_table="PerAtomRMSE"
--default_dtype="float64"
--restart_latest
--device=cuda
--seed=1615767616
--energy_key="energy"
--forces_key="forces"
--energy_weight=5.0
--forces_weight=10.0
--ema
--ema_decay=0.95
--swa
--swa_lr=0.001
--swa_energy_weight=10.0
--swa_forces_weight=5.0

Could someone help me with this warning? Thank you^^

[ilyes319]

Hey, what is your mace version?

[YaoqiPang]

mace-torch 0.3.10.

With setting above, the middle layer is torch.Size([98304]), however mace use [128, 3, 128] to match torch.Size([98304]).

[ilyes319]

@YaoqiPang Can you email me your model here: [email protected] so I can debug this