数据参数设置

数据集路径

关键词 类型 默认值 描述
db_path str ./ 包含你构建的数据集的文件夹
split str split 分割训练和验证集的 JSON 文件名(不包括 .json 后缀)

有关数据集的构建和存放方式可以参考数据集构建。不过还请注意的是,尽管 LMDB 支持并发读取,但是多个训练任务使用同一个数据集的话,还是尽量复制多分分别去读以防进程冲突,有条件的话最好能放到节点专属的固态硬盘上。

截断半径

关键词 类型 默认值 描述
cutoff float 5.0 截断半径数值

关于截断半径的介绍可以参考图的建立。单位与模型超参数中设置的默认长度单位一致。

当然你会发现在定义模型超参数的时候也需要指定截断半径,这两者的区别是,模型超参数的截断半径定义了诸如包络函数等模块的截断半径值;而这里数据集定义的截断半径是预处理数据形成分子图时所使用的截断半径。所以从定义上来说这里的截断半径大一点其实不要紧,因为超过的部分会被处理为贡献为 0,但是保持一致终归是没有错的。

训练标签

关键词 类型 默认值 描述
targets list[str] - 用于训练的标签的名字
base_targes list[str] null 用于 Δ-ML 的标签,注意和 targets 一一对应
default_dtype str float32 使用的浮点数精度

这个还是直接举例说明:

data:
  ...
  targets: [energy, forces]
  base_targets: [base_energy, base_forces]
  ...

意思是能量和力都是用 Δ-ML,学习的目标就是上下两个标签对应的值的差值。当然数据集里必须有这几个声明了的标签。

提到 Δ-ML,有些要注意的是,由于不同的方法(泛函和基组等)对应的能量量级有些不同,比如全电子 DFT 和使用赝势的半经验 xTB 方法给出的能量量级就完全不同,直接相减显然是不明智的。建议参考LMDB 数据集制作部分给出的小贴士扣除原子能量以保持量级一致。

标签归一化

关键词 类型 默认值 描述
node_scale bool / float false 标签的每原子标准差,将输出乘以这个标准差。如果是 true 则在线根据标签计算标准差;如果为浮点数,则直接使用这个值
node_shift bool / float false 标签的每原子平均值,将输出加上这个均值。如果是 true 则在线根据标签计算均值;如果为浮点数,则直接使用这个值
max_num_sample int 10000000 用于在线计算的平均值和标准差的最大样本数量

标签归一化的作用是可以提升模型训练效率。这里使用 Z-Score 标准化:

\[ y^\prime_i = \frac{y_i - \mu}{\sigma} \]

这里 \(\mu\) 是标签的均值,\(\sigma\) 是标签的标准差。由于能量大小和原子数量也有关系,所以这里我们使用每原子的能量计算均值和标准差,即:

\[ \mu_a = \frac{1}{N} \sum_i^N \frac{y_i}{n_i} \\ \sigma_a = \sqrt{ \frac{1}{N} \sum_i^N \left( \frac{y_i}{n_i} - \mu_a \right) } \\ y^\prime_i = \frac{y_i - n_i \mu_a}{n_i \sigma_a} \]

比起修改标签归一化,在模型的每原子能量的输出上,乘以标准差,加上均值这个操作更便利一些,即:

\[ E_i^\prime = \left( E_i \cdot \sigma_a \right) + \mu_a \]

实际上程序用了更讨巧的方法,就是直接拿出最后的 Linear 层,初始化的时候直接在 Weight 上乘上 \(\sigma_a\),在 Bias 上加上 \(\mu_a\)(遵循奥卡姆剃刀原则)。

批处理大小

关键词 类型 默认值 描述
batch_size int 64 训练集批次大小
valid_batch_size int 64 验证集批次大小

训练集和验证集的大小,会平均分配到每个 GPU 上。也就是每张 GPU 上的实际批次大小是 batch_size // n_gpu