MACE多头训练与模型拆分

下载本专题输入、Si实算数据、教学模型与工具（ZIP）。本例让同一个Si构型同时拥有PBE与PBEsol两套DFT真值，用两个读出头区分计算级别。两套标签都来自实际VASP计算，未用缩放或常数平移伪造第二套数据。

1 准备两套数据#

module load mace/0.3.16
python inspect_data.py data/PBE-train.xyz
python inspect_data.py data/PBEsol-train.xyz

两套数据各32训练、8验证、8测试；结构编号的拆分一致。不能让一个结构在一个头进入训练、另一个头进入测试，否则测试集不再独立。

2 多头配置与训练#

cat configs/multihead64.yml
bash train.sh multihead64

基础参数与从头训练相同，额外配置：

heads:
  PBE:
    train_file: ../../data/PBE-train.xyz
    valid_file: ../../data/PBE-valid.xyz
    E0s: average
    energy_key: REF_energy
    forces_key: REF_forces
    stress_key: REF_stress
    virials_key: REF_virials
  PBEsol:
    train_file: ../../data/PBEsol-train.xyz
    valid_file: ../../data/PBEsol-valid.xyz
    E0s: average
    energy_key: REF_energy
    forces_key: REF_forces
    stress_key: REF_stress
    virials_key: REF_virials

两个头共享大部分表示，各自学习能量读出。这里是多DFT级别联合训练，不是把PBEsol当作第二个元素，也不是让用户在推断时混合两种能量。官方多头训练说明了这个结构。

3 推断时明确选头#

from mace.calculators import MACECalculator
calc = MACECalculator(model_paths="models/multihead64.model",
                      device="cpu", default_dtype="float64", head="PBEsol")
atoms.calc = calc
print(atoms.get_potential_energy())

用PBEsol头的预测与PBEsol测试真值比较。实际两个头的结果为：

4 拆出单个头#

bash calculate.sh multihead
ls work/multihead-*.model

助手调用原生工具：

mace_select_head models/multihead64.model --head_name PBE   --target_device cpu --output_file work/multihead-PBE.model

PBEsol同理。拆分前后，能量、力、应力和virial逐结构对照，最大差均低于 1E-8 的对应原始单位。附件同时保留完整双头和已拆分的CPU模型。

mace_eval_configs --configs data/PBEsol-test.xyz   --model models/multihead64.PBEsol.model --head PBEsol   --output work/PBEsol-head.xyz --device cpu --default_dtype float64   --batch_size 1 --compute_stress
python inspect_data.py work/PBEsol-head.xyz

拆头是部署操作，不会重新训练。改变头名不会把一套DFT标签自动转换为另一套；新计算级别需要对应数据。