FAIRChem 自定义数据集构建指南使用 ASE 数据库与 ASE 可读文件训练与推理【免费下载链接】ocpFAIR Chemistrys library of machine learning methods for chemistry项目地址: https://gitcode.com/GitHub_Trending/oc/ocpFAIRChem 的模型如 UMA、eSCAIP、AllSCAIP 等接受任意化学结构作为输入预测能量energy、力forces、位置positions与应力stress。本文以 docs/core/common_tasks/ase_dataset_creation.md 为核心系统讲解如何在 FAIRChem 中构建自定义数据集既可以直接读取已有的ASE 数据库ASE Database也可以从ASE 可读的各类结构文件CIF、POSCAR、traj、XYZ 等直接训练与推理并深入到 ase_datasets.py 的源码实现与测试用例帮助你掌握每种dataset.format的配置细节、性能取舍与完整落地流程。一、FAIRChem 中的数据集体系fairchem提供面向「以任意化学结构为输入、预测能量/力/位置/应力」任务的训练与评估代码可以当作研究项目的基础脚手架使用。在正式进入自定义数据集之前先了解 FAIRChem 生态中已经内置的公开数据集与对应任务有助于理解数据集接口的定位。仓库文档对各类数据集的说明位于 docs 目录OC20Open Catalyst 2020催化剂表面吸附能量/力基准OC22扩展的氧化物催化数据集ODAC23直接空气捕集DAC材料数据集OC20DenseOC20 的稠密采样版本OC20NEB用于 NEBNudged Elastic Band路径任务OMat24无机材料数据集OMol25有机分子数据集OMC25分子催化数据集。在 OC20、OC22 等内置数据集之外FAIRChem 提供多种途径在自有数据上训练与评估模型。原文档明确指出将数据写成 LMDB 是性能最优的选择而ASE 系数据集格式ASE DB、ASE 可读文件是为已有数据、希望快速试用 fairchem 工具链、又不想学习 LMDB 细节的用户提供的便捷方案。从源码角度看所有 ASE 系数据集都继承自 ase_datasets.py 中的抽象基类AseAtomsDataset它把「读取ase.Atoms对象」与「转换为图数据对象AtomicData」两个阶段解耦子类只需实现get_atoms(idx)按标识符返回ase.Atoms与_load_dataset_get_ids(config)初始化并返回全部样本标识符列表两个方法图转换统一通过AtomicData.from_ase见 atomic_data.py完成a2g_args中的r_energy/r_forces/r_stress等参数即透传给该构造器。目前注册在fairchem.core.datasets中的 ASE 系数据集共有三个与配置文件中的dataset.format一一对应format值数据集类适用场景ase_dbAseDBDataset数据已存在 ASE Database 中JSON / SQLite / 数据库服务器后端ase_readAseReadDataset每个文件只包含一个结构如单个 CIF、POSCARase_read_multiAseReadMultiStructureDataset每个文件包含多个结构如 ASE.traj、多帧 XYZ二、直接使用自定义 ASE 数据库format: ase_db如果你的数据已经存在于ASE Database中那么在运行训练/推理之前不需要任何额外预处理直接在配置中把dataset.format从 LMDB 切换为ase_db即可。吞吐量提示ASE DB 后端虽然吞吐量可能不及 LMDB但对「12 块 GPU 训练规模适中的数据集」或「单 GPU 推理」来说通常已经足够快。若要更有效地利用更多资源建议把数据写成 LMDB。原文档给出的ase_db完整配置如下dataset: format: ase_db train: src: # The path/address to your ASE DB connect_args: # Keyword arguments for ase.db.connect() select_args: # Keyword arguments for ase.db.select() # These can be used to query/filter the ASE DB a2g_args: r_energy: True r_forces: True # Set these if you want to train on energy/forces # Energy/force information must be in the ASE DB! keep_in_memory: False # Keeping the dataset in memory reduces random reads and is extremely fast, but this is only feasible for relatively small datasets! include_relaxed_energy: False # Read the last structures energy and save as y_relaxed for IS2RE-Direct training val: src: a2g_args: r_energy: True r_forces: True test: src: a2g_args: r_energy: False r_forces: False # It is not necessary to have energy or forces if you are just making predictions.2.1 srcASE DB 的多种指定方式AseDBDataset._load_dataset_get_idsase_datasets.py对src的处理非常灵活支持以下形式单个文件路径如src: /path/to/my_asedb.db单个文件夹src指向目录时目录下的每个文件都会被尝试当作 ASE DB 连接无法干净连接的会记录 debug 日志并跳过路径列表 / 文件夹列表src为 YAML 列表时逐个处理其中的文件或目录并拼接为一个大数据集glob 字符串如果src既不是文件也不是目录会被当作 glob 模式匹配出一批 ASE DB 文件。连接时通过connect_dbase_datasets.py调用ase.db.connect(address, **connect_args)特别地若地址中包含aselmdb如.aselmdb扩展名会自动设置readonlyTrue与use_lock_fileFalse以保证只读安全。2.2 connect_args / select_args连接与查询过滤connect_args是ase.db.connect()的关键字参数用于控制数据库连接方式例如 SQLite 的只读模式、后端类型等select_args是ase.db.select()的关键字参数用于查询/过滤数据库例如按 key-value 条件只选取部分结构作为训练集。源码层面ID 列表的获取有一个性能优化如果数据库对象本身暴露了ids属性且select_args为空则直接复用db.ids否则必须遍历db.select(**select_args)逐条收集 ID——后者在大数据集上很慢。get_atoms则通过bisect在多个 DB 之间定位样本并把 DB 行中的data字典写回atoms.info供后续图转换读取。2.3 a2g_args控制读取哪些物理量a2g_args透传给AtomicData.from_aseatomic_data.py决定将原子的哪些属性写入图数据对象r_energy: True/r_forces: True训练能量/力时开启能量与力信息必须存在于 ASE DB 中例如通过SinglePointCalculator写入或在atoms.info中保存energy/forces键r_stress: True需要训练应力时开启from_ase会处理 6 分量 Voigt 记法、3×3 矩阵与 9 分量展开等多种应力存储形式见 atomic_data.pyinclude_relaxed_energy: True读取结构中最后一个弛豫终态结构的能量保存为energy_relaxed字段用于 IS2RE-Direct 类任务的训练该参数不要求物理量在 DB 里显式命名而是由各数据集类自行实现读取逻辑。从当前源码看from_ase签名中r_energy已被标记为 deprecated保留仅为兼容能量/力在计算器结果存在时即被读取r_stress则实际控制应力是否读取——配置层面保持原文档推荐的显式写法即可。2.4 keep_in_memory小数据集的加速开关keep_in_memory: True会把整个数据集缓存在 CPU 内存中避免每次随机读取磁盘能显著加快训练速度。前提是数据集足够小、能放进内存——源码通过functools.cache装饰__getitem__实现ase_datasets.py。对于需要多轮迭代epoch的小型数据集尤其推荐大数据集则不建议开启。三、仓库中的真实 ase_db 配置范例ase_db并非仅在文档中存在仓库的真实配置大量使用该格式。以 configs/escaip/training/dataset/fair_cluster_mptrj.yaml 为例可以看到生产级配置如何组合a2g_args、key_mapping与transformsmptrj_train: splits: train: src: - data/mptrj/train/train_mptrj.aselmdb format: ase_db a2g_args: r_energy: True r_forces: True r_stress: ${regress_stress} key_mapping: energy: mptrj_energy forces: ${mptrj_forces_key} stress: ${mptrj_stress_key} transforms: common_transform: dataset_name: mptrj stress_reshape_transform: dataset_name: mptrj mptrj_val: splits: val: src: - data/mptrj/val/val_mptrj.aselmdb format: ase_db a2g_args: r_energy: True r_forces: True r_stress: ${regress_stress} key_mapping: energy: mptrj_energy forces: ${mptrj_forces_key} stress: ${mptrj_stress_key} transforms: common_transform: dataset_name: mptrj stress_reshape_transform: dataset_name: mptrj要点解读src以列表形式给出.aselmdb文件触发上文提到的aselmdb只读安全开关key_mapping将数据集中属性名映射为模型期望的名称源码在AseAtomsDataset.__getitem__中通过rename_data_object_keys实现见 ase_datasets.py${regress_stress}等占位符由 Hydra 插值语法注入transforms支持在数据集对象上施加预处理如common_transform、stress_reshape_transform由DataTransforms统一调度。同样的format: ase_db用法还出现在 configs/allscaip/dataset/omol.yaml、configs/uma/training_release/dataset/uma.yaml 等大量配置中可作为不同模型族下自定义数据集的参考模板。四、直接在 ASE 可读文件上训练/推理format: ase_read / ase_read_multiFAIRChem 支持直接读取 ASE 可读的文件进行训练与推理无需任何预处理。需要特别注意的是警告这种方式只推荐用于较小的数据集——大量小文件组成的目录在多数计算基础设施上无法高效扩展随机 IO 会成为瓶颈。AseReadDataset与AseReadMultiStructureDataset使用ase.io.read读取文件支持的格式覆盖 ASE 文档中列出的全部可读类型CIF、POSCAR、traj、XYZ、OUTCAR 等。两种格式按「每个文件包含的结构数」区分。4.1 单结构文件format: ase_read该格式假定每个文件恰好包含一个结构dataset: format: ase_read train: src: # The folder that contains ASE-readable files pattern: # Pattern matching each file you want to read (e.g. */POSCAR). Search recursively with two wildcards: **/*.cif. include_relaxed_energy: False # Read the last structures energy and save as y_relaxed for IS2RE-Direct training ase_read_args: # Keyword arguments for ase.io.read() a2g_args: # Include energy and forces for training purposes # If True, the energy/forces must be readable from the file (ex. OUTCAR) r_energy: True r_forces: True keep_in_memory: False参数说明对照 ase_datasets.py 的实现src包含 ASE 可读文件的文件夹路径。若指向文件_load_dataset_get_ids会直接抛出ValueErrorpattern匹配待读取文件的 glob 模式如*/POSCAR、*.cif、*.xyz使用两个通配符可递归搜索如**/*.cif。实现通过Path.glob完成默认*ase_read_args透传给ase.io.read()的关键字参数。源码特别检查了index参数如果index中包含:切片语法会抛出NotImplementedError提示多结构文件请使用ase_read_multia2g_args同上节训练时需要能量/力则开启且能量/力必须能从文件中读出例如从 OUTCAR 解析include_relaxed_energy开启后实现会把ase_read_args的index强制改为-1读取文件最后一帧作为弛豫终态并调用atoms.get_potential_energy(apply_constraintFalse)得到弛豫能量keep_in_memory同前小数据集推荐True。另外AseAtomsDataset默认会通过apply_one_tagsase_datasets.py为完全没有 tag 的结构打上全 1 的 tag。原因是部分模型如 GemNet-OC默认只对非零 tag 的原子计算三体/四体相互作用若结构无 tag 会直接报错。若你的结构自带 tag 语义skip_if_nonzero: True默认或想完全关闭skip_always: True可通过atoms_transform_args调整。4.2 多结构文件format: ase_read_multi该格式支持读取每个文件包含多个结构的数据典型如 ASE.traj文件、多帧 XYZdataset: format: ase_read_multi train: index_file: Filepath to an index file which contains each filename and the number of structures in each file. e.g.: /path/to/relaxation1.traj 200 /path/to/relaxation2.traj 150 ... # If using an index file, the src and pattern are not necessary src: # The folder that contains ASE-readable files pattern: # Pattern matching each file you want to read (e.g. *.traj). Search recursively with two wildcards: **/*.xyz. ase_read_args: # Keyword arguments for ase.io.read() a2g_args: # Include energy and forces for training purposes r_energy: True r_forces: True keep_in_memory: False强烈建议提供index_file。原因从源码可以看得很清楚ase_datasets.py提供index_file时数据集启动时只读取索引文件每行文件名 结构数并按行展开出(filename, frame_index)的样本 ID 列表src与pattern可省略不提供index_file时数据集被迫在启动阶段把所有文件全部读一遍来统计结构数量这是大数据集上非常显著的一次性开销。索引文件格式为每行路径 数量/path/to/relaxation1.traj 200 /path/to/relaxation2.traj 150其余参数ase_read_args、a2g_args、keep_in_memory语义与ase_read一致实现上会为ase_read_args设置默认index: :以读取全部帧。include_relaxed_energy同样取文件最后一帧作为弛豫终态get_relaxed_energy读取[-1]帧。读取文件时还可通过use_tqdm默认True控制启动阶段的进度条显示。五、性能取舍与工程实践建议5.1 三种格式的适用边界场景推荐格式理由数据已入库、规模中等、12 卡训练ase_db零预处理、支持查询过滤吞吐对中小规模足够少量散文件、快速试跑ase_read/ase_read_multi免预处理但大量小文件 IO 不具扩展性大规模训练、多卡并行LMDBlmdb格式吞吐最高能有效利用更多资源5.2 数据集与训练流程的集成方式ASE 数据集通过统一的create_dataset工厂接入训练框架base_dataset.py函数依据配置中的dataset.format从 registry 解析出数据集类默认lmdb再处理splits划分、seed、max_atoms按原子数过滤、subset_to按元数据键过滤、first_n/sample_n/no_shuffle三者互斥等数据集级配置最终返回 PyTorchSubset。这意味着你只需要在 Hydra 配置中写清楚format: ase_db或ase_read/ase_read_multi及对应参数训练、评估、推理链路即可无缝复用。具体启动训练的命令遵循 FAIRChem v2 的统一 CLI 约定详见 训练指南CLI 只接受一个参数——Hydra YAML 路径其余一切配置均通过配置文件与 Hydra override 语法提供。例如把你的自定义数据集配置放到自己的配置组后可以用类似下面的命令启动fairchem -c path/to/your_config.yaml如果需要临时覆盖数据集路径或开关可使用 Hydra 的 override 语法如datasettrain src/path/to/your_ase_db具体覆盖方式以你的配置组结构为准。5.3 测试用例佐证仓库中的 tests/core/datasets/test_ase_datasets.py 对三种 ASE 数据集均有覆盖可作为自定义数据集正确性的行为参考test_ase_dataset对AseDBDataset的单文件、文件夹、路径列表、.aselmdb四种src形态逐一验证样本数、forces形状(natoms, 3)、stress形状(1, 3, 3)与sid注入test_ase_read_dataset把结构写成多个.cif后验证AseReadDataset的src pattern读取test_ase_multiread_dataset写入多帧.traj同时验证index_file路径与无索引路径并断言include_relaxed_energy下energy_relaxed的正确性首帧不等于首帧能量、末帧等于末帧能量test_db_add_delete验证 ASE DB 中结构增删后数据集长度随之变化test_empty_datasetsrc下无任何有效数据时会抛出ValueError对应 ase_datasets.py 的空数据集校验。六、小结FAIRChem 为自定义数据提供了「零预处理」的 ASE 系数据集方案覆盖三种典型形态ase_db直接连接既有 ASE Databasesrc支持文件/目录/列表/glob配合connect_args、select_args完成连接与过滤a2g_args控制能量/力/应力的读取ase_read每个文件一个结构用src pattern含**递归定位文件适合小规模试跑ase_read_multi每个文件多个结构务必提供index_file避免启动时全量扫描适合轨迹类数据。无论选择哪种格式keep_in_memory: True都是小数据集显著提速的利器而当数据规模增长到需要充分利用更多 GPU 时则应考虑迁移到 LMDB 格式。通过源码ase_datasets.py、base_dataset.py、atomic_data.py、真实配置如 fair_cluster_mptrj.yaml与测试用例test_ase_datasets.py的相互印证你可以在完全理解底层机制的前提下把自己的结构数据快速接入 FAIRChem 的训练与推理管线。【免费下载链接】ocpFAIR Chemistrys library of machine learning methods for chemistry项目地址: https://gitcode.com/GitHub_Trending/oc/ocp创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考