PaddleMaterials Trainer 文档

1.1 快速上手

# 性质预测（单卡）
python property_prediction/train.py -c property_prediction/configs/spherenet/spherenet_qm9_U0.yaml

# 原子间势（多卡）
python -m paddle.distributed.launch --gpus="0,1,2,3" \
    interatomic_potentials/train.py -c interatomic_potentials/configs/chgnet/chgnet_mptrj.yaml

# 结构生成（单卡）
python structure_generation/train.py -c structure_generation/configs/diffcsp/diffcsp_mp20.yaml

# 临时覆盖配置项（命令行追加 Key=Value）
python <task>/train.py -c <config.yaml> Trainer.max_epochs=50

1.2 我想做哪件事

1.3 配置文件

Global:
  do_train / do_eval / do_test : True
Trainer:
  max_epochs / max_iter / seed / output_dir
  eval_freq | eval_freq_steps     # 二选一（epoch 级 vs 步级 eval，互斥）
  save_freq | save_freq_steps     # 检查点节奏
  pretrained_model_path / resume_from_checkpoint
  use_amp / eval_with_no_grad / gradient_accumulation_steps
  best_metric_indicator / name_for_best_metric / greater_is_better
  use_tensorboard / use_visualdl / use_wandb
Model / Dataset / Optimizer / Vocabulary / Metric  # 按 __class_name__ 注册

1.4 Trainer 配置项

1.5 常见任务

从头训练

python <task>/train.py -c <config.yaml>

断点续训

python <task>/train.py -c <config.yaml> Trainer.resume_from_checkpoint=./output/.../latest.pdparams

加载预训练微调

python <task>/train.py -c <config.yaml> Trainer.pretrained_model_path=./pretrained/model.pdparams

只评估 / 测试

python <task>/train.py -c <config.yaml> \
    Global.do_train=False Global.do_eval=True \
    Trainer.pretrained_model_path=./output/.../best.pdparams

步级 eval（长 epoch 用）

Trainer.eval_freq_steps=200 Trainer.save_freq_steps=2000

1.6 运行后产物

output_dir 运行时自动加 _t_<时间戳>_s_<seed> 后缀。

1.7 训练时间线

2.1 装配全景

2.2 注册约定

Model:
  __class_name__: MyModel        # 类名，需 import 进对应 __init__.py
  __init_params__: { hidden_dim: 128 }

2.3 模型契约

class MyModel(paddle.nn.Layer):
    def __init__(self, hidden_dim=128, vocab=None, **kwargs):
        super().__init__(); self.vocab=vocab; self.target_name="energy"
    def forward(self, batch):
        pred = self._compute(batch)
        loss = some_loss(pred, batch["label"])
        return {"loss_dict": {"loss": loss, "mae": paddle.mean(paddle.abs(pred-batch["label"]))},
                "pred_dict": {self.target_name: pred}}
    def before_train(self, trainer): ...   # 可选 hook

loss_dict["loss"] 是反向唯一入口；其它键可作 best 指标（name_for_best_metric）。

生成式模型（扩散：DiffNMR / DiffCSP / MatterGen）

生成式模型的 训练 也返回 loss_dict（BaseTrainer 优化），但 生成/采样 走独立的 sample() 方法 + Sampler，不经 Trainer：

# 训练（BaseTrainer 优化）
def forward(self, batch):
    return {"loss_dict": {"loss": total_loss, "loss_coord": ..., "loss_atom": ...}}

# 生成（sample.py 入口 + MolecularSampler/StructureSampler，不经 Trainer）
def sample(self, batch, num_inference_steps=1000, **kwargs):
    return generated_structures

2.4 数据集契约

class MyDataset(paddle.io.Dataset):
    def __getitem__(self, idx):
        return {"graph": pgl_graph, "x": np.ndarray, "y": np.ndarray, "id": idx, "info": {...}}
    def __len__(self): ...

2.5 指标

(a) 函数式：def metric(pred,label)->scalar，注册到 Metric 段。
(b) 流式：实现 update_step/compute_epoch/reset，框架在 epoch 边界聚合。
best 口径：best_metric_indicator + name_for_best_metric + greater_is_better。

2.6 训练入口

config = OmegaConf.to_container(OmegaConf.merge(OmegaConf.load(c), OmegaConf.from_dotlist(args)), resolve=True)
vocab = build_vocab(config.get("Vocabulary"))
model = build_model(config["Model"], vocab=vocab)
loaders = read_dataloader_config(config, vocab=vocab)
opt, lr = build_optimizer(config["Optimizer"], model, config["Trainer"]["max_epochs"], len(loaders[0]))
trainer = BaseTrainer(config["Trainer"], model, *loaders[:2], optimizer=opt, lr_scheduler=lr,
                      compute_metric_func_dict=build_metric(config.get("Metric")))
if do_train: trainer.train()
if do_eval:  trainer.eval(loaders[1])
if do_test:  trainer.eval(loaders[2])

2.7 BaseTrainer 能力

2.8 最小骨架

# my_model.py
class MyModel(paddle.nn.Layer):
    def __init__(self, hidden_dim=64, vocab=None, **kw):
        super().__init__(); self.linear=paddle.nn.Linear(hidden_dim,1)
    def forward(self, batch):
        pred=self.linear(batch["x"]).squeeze(-1)
        loss=paddle.nn.functional.mse_loss(pred,batch["y"])
        return {"loss_dict":{"loss":loss,"mae":paddle.mean(paddle.abs(pred-batch["y"]))},
                "pred_dict":{"y":pred}}
# ppmat/models/__init__.py 加: from .my_model import MyModel
# config.yaml: Model/Dataset/Optimizer/Trainer/Global 同标准结构

2.9 常见陷阱

• 类必须 import 进 __init__.py，否则 eval(__class_name__) 报 NameError。
• loss_dict 必须有 "loss"。
• collate 要与样本类型匹配（pgl.Graph / Data 各有路径）。
• 分布式自定义存盘要尊重 rank（只 rank0）。
• vocab 仅当构造函数有该参数才注入。
• ${} 插值在 resolve=True 展开。

3.1 代码地图

3.2 架构分层

分层原则：入口只编排；Trainer 不感知具体模型（只认 forward 契约）；组件层走注册反射；基础层可独立单测。

3.3 新增一个任务（checklist）

1. 建目录 <task>/，含 train.py / predict.py(或 sample.py) / configs/。
2. 实现模型（§2.3 契约），import 进 ppmat/models/__init__.py。
3. 实现数据集（§2.4），选/写 collate_fn，import 进 ppmat/datasets/__init__.py。
4. （按需）指标、graph converter、词表（注册到 vocab）。
5. config 模板：Global/Trainer/Model/Dataset/Optimizer/Vocabulary/Metric。
6. README：命令 + 结果表 + 数据准备说明。
7. （可选）预训练权重：上传 BCE → MODEL_REGISTRY 登记 → 按 §3.8 打包。
8. （可选）推理：实现/复用 ppmat/predictor/ 子类。
9. 测试：test/test_<task>.py（契约、collate、小样本前向）。
10. 冒烟：max_epochs=1 跑通 → 提交。

3.4 Trainer 内部机制

方法职责边界

两套 eval 节奏（互斥）

检查点与分布式

• 前缀：best（_determine_best_metric）、latest（save_freq / 步级 eval）、step_N（save_freq_steps）。
• save_load.save_checkpoint 仅 rank0 写；多卡 cache 构建后 dist.barrier() 同步。
• 步级 step_N 检查点在 batch 循环顶层独立判断（不依赖 eval 命中）。

lr / AMP / 累积

• ReduceOnPlateau 按 indicator/indicator_name 取值 step；其他调度器直接 step。
• AMP：paddle.amp.GradScaler + amp.decorate；梯度累积按 gradient_accumulation_steps。

3.5 缓存子系统（密度数据集）

• 首次构建把 field/graph 序列化到 <cache>/fields|graphs/{idx:010d}.pkl；后续直读。
• 校验：比 build_field_cfg.pkl / build_graph_cfg.pkl / graph_vocab.pkl（is_equal）。
• 多进程：ProcessPoolExecutor(spawn) + worker initializer，避免继承 CUDA 状态。
• MD17 子类：FFT→CUBE 物化（首次），复用 base cache（见 §Ⅰ/Ⅱ 的数据契约）。

3.6 配置系统

• OmegaConf.load → merge(dotlist) → to_container(resolve=True)。
• ${a.b} 插值在 resolve 阶段展开；CLI 用 dotlist A.b.c=v。
• rank0 把最终 config 存 <output_dir>/<name>.yaml 快照。
• 未知键：构造函数用 **kwargs 吞掉或 strict_unused 报错（build_model）。

3.7 下载与注册表

• download.get_datasets_path_from_url(url, md5) / get_weights_path_from_url：下到 DATASETS_HOME/WEIGHTS_HOME，md5 校验，自动解压。
• MODEL_REGISTRY（models/__init__.py）：name → BCE 权重 url；VOCAB_REGISTRY（vocab）：name → 词表 url + md5。
• 数据集子类用类属性 url/md5 声明下载源；split_url/split_md5 为预留字段。

3.8 预训练模型包契约

注册的预训练模型须用固定 archive 布局（见 docs/model_package.md）：

<model_name>.zip
└── <model_name>/
    ├── <model_name>.yaml      # 配置
    └── checkpoints/
        └── best.pdparams       # 权重

包名、顶层目录、配置文件名须与 MODEL_REGISTRY 的 key 一致。FieldPredictor 等按此布局解析。

3.9 测试

• test/：pytest，覆盖 spherenet/scatter/md17_dataset/download 等基础路径。
• 建议补：collate 行为、cache 命中/未命中、模型 forward 契约（loss_dict/pred_dict）。
• 运行：python -m pytest test/ -q（注意某些环境对 mmap/网络挂载敏感）。
• Cython 扩展（如 mattersim）需先 setup.py build_ext --inplace 再跑相关测试。

3.10 依赖与发布

• requirements.txt：运行时依赖（setup.py 读取并去重）。pyproject.toml：构建/开发依赖。
• setup.py：Cython 扩展编译 + setuptools-scm 版本。
• .pre-commit-config.yaml：black / ruff / eof / yaml / 大小写等钩子；提交前 pre-commit run。
• 新增运行时 import 的三方库（如 plotly/matplotlib/kaleido）须同时进 requirements。

3.11 代码风格约定

• 节点特征键名：cart_coords / atom_types / x（勿再用旧 pos/atomic_number）。
• 模型统一 forward→{loss_dict,pred_dict} 契约。
• 组件走 __class_name__/__init_params__ 注册；类 import 进对应 __init__.py。
• 日志用 ppmat.utils.logger；分布式写盘只 rank0。
• 用户输入校验抛 ValueError/TypeError（勿用 assert）。
• 遵循 black/ruff；行宽 88。

3.12 维护注意事项（演进时）

• 改 forward 契约 / 节点特征键名 / 注册模式属破坏性变更，需同步所有模型与测试。
• 改缓存布局或校验键会失效旧缓存——升 schema_version 或文档说明清缓存。
• 改 build_* 签名注意向后兼容（vocab 等可选注入）。
• eval 节奏 / 检查点时序的改动要有测试覆盖（易出静默不存盘/指标污染）。
• 预训练权重 / 词表 URL 改动要同步 MODEL_REGISTRY / VOCAB_REGISTRY 的 md5。
• 新增依赖同步 requirements + 文档；重模型/数据 import 避免放到包 __init__ 顶层（影响 import 时延）。