DataFlex注册表系统详解：三步注册加载你的新数据调度算法

DataFlex注册表系统详解：三步注册加载你的新数据调度算法

【免费下载链接】DataFlex可用于大模型训练时动态进行训练动态训练数据选择、领域比例调整及动态加权，提升训练速度和性能，与 LLaMA-Factory 无缝集成，提供灵活强大的训练控制能力。项目地址: https://gitcode.com/OpenDCAI/DataFlex

DataFlex 是一个用于大模型训练的动态数据调度工具，支持训练数据选择、领域比例调整与动态加权，并与 LLaMA-Factory 无缝集成。它的核心是一个轻量注册表系统：所有算法组件（Selector / Mixer / Weighter / Reorder）都通过装饰器注册、按名称从 YAML 配置加载。本文带你 3 步注册并加载你自己的数据调度算法。🚀

DataFlex 注册表是什么？

注册表位于 registry.py，本质上是一个「类型 → 名称 → 类」的两层字典：

REGISTRY = Registry() def register_selector(name): return REGISTRY.register("selector", name) def register_mixer(name): return REGISTRY.register("mixer", name) def register_weighter(name): return REGISTRY.register("weighter", name) def register_reorder(name): return REGISTRY.register("reorder", name)

它只做三件事：

build是新手最该关注的一点：它会检查组件__init__的函数签名，只传入类真正声明的参数，多余的配置会被自动过滤——你不需要手写任何参数解析代码。

三步注册你的新算法（以 Selector 为例）

完整的官方扩展指南见 how_to_add_algorithm.md，核心流程如下。

第一步：继承基类并加装饰器

新建一个文件（如src/dataflex/train/selector/my_selector.py），继承 base_selector.py 中的Selector基类，并用装饰器注册：

from dataflex.core.registry import register_selector from .base_selector import Selector @register_selector("my_method") class MySelector(Selector): def __init__(self, dataset, accelerator, data_collator, cache_dir, my_param=0.5): super().__init__(dataset, accelerator, data_collator, cache_dir) self.my_param = my_param def select(self, model, step_id: int, num_samples: int, **kwargs): # 在此实现你的选样逻辑，返回样本索引列表 return list(range(num_samples))

💡 基类已提供warmup（随机预热采样）、候选池管理等能力，子类只需实现select一个抽象方法。

第二步：注册导入

在 selector/init.py 中加一行导入，让注册表能"看到"你的类：

from .my_selector import *

Mixer、Weighter、Reorder 同理，分别在对应的train/mixer/、train/weighter/、train/reorder/包中导入即可。

第三步：在 components.yaml 中加一段配置

打开 components.yaml，在selectors:分组下添加你的组件与默认参数：

selectors: my_method: name: my_method params: my_param: 0.8 cache_dir: ../dataflex_saves/my_method_output

训练时如何按名称加载组件？

训练配置中只需写三个字段：

train_type: dynamic_select components_cfg_file: src/dataflex/configs/components.yaml component_name: my_method

对应的 Trainer（如 select_trainer.py）会按流程完成加载：

1. load_component读取 components.yaml 中my_method的参数（load_component.py）；
2. 调用REGISTRY.build("selector", "my_method", runtime=..., cfg=...)拿到组件实例；
3. 运行时参数（dataset、accelerator 等）优先于 YAML 参数，缺失项自动补齐。

之后 Trainer 会在 warmup 之后、每隔update_step步自动调用你的select()/mix()/get_weighted_loss()，无需额外代码。

相关文件速查表

常见问题

Q：为什么注册后component_name找不到组件？A：多半是第二步漏了导入——装饰器只在模块被导入时才执行，务必确认对应包的__init__.py已引用你的文件。

Q：我的组件参数和 Trainer 传入的参数冲突了？A：build中运行时参数（runtime）优先于 YAML 参数（cfg），这是刻意设计，保证训练状态实时生效。

Q：注册表会重复注册报错吗？A：会。同名重复注册会抛出ValueError，方便你在开发阶段尽早发现命名冲突。

掌握这套「装饰器注册 → 包内导入 → YAML 配置」三步流程后，你就可以像插件一样持续向 DataFlex 扩展自己的数据调度策略了。✅

【免费下载链接】DataFlex可用于大模型训练时动态进行训练动态训练数据选择、领域比例调整及动态加权，提升训练速度和性能，与 LLaMA-Factory 无缝集成，提供灵活强大的训练控制能力。项目地址: https://gitcode.com/OpenDCAI/DataFlex