跳到内容

添加数据集

Ludwig 数据集动物园 (Ludwig Dataset Zoo) 收集了网络上各种数据集,并方便地内置于 Ludwig 中。

Ludwig 数据集自动化了从 Kaggle 等网站下载数据时的凭据管理、将多个文件合并为一个数据集、共享数据解析代码以及将数据集直接加载到可用于训练 Ludwig 模型的数据框中。

  • Ludwig Datasets API 位于 ludwig/datasets/ 中。
  • 数据集配置定义在 ludwig/datasets/configs/ 下。
  • 特定数据集的自定义加载器位于 ludwig/datasets/loaders/ 中。

通过提供数据集配置 .yaml 文件,即可在 Ludwig 中使用数据集。对于许多数据集而言,创建此 YAML 文件是唯一必要的步骤。

1. 创建新的数据集配置

ludwig/datasets/configs/ 下创建一个新的 .yaml 文件,文件名为数据集名称。此配置文件必须包含以下必需的键:

  • version: 数据集的版本
  • name: 数据集的名称。这是将用于导入或传递给 get_datasets(dataset_name) 的名称。
  • description: 数据集的人类可读描述。可能包含多行文本和链接。
  • download_urlskaggle_competitionkaggle_dataset_id 之一。

支持的压缩存档和数据文件类型将根据文件扩展名自动推断。

注意

如果数据集具有在研究论文或 Kaggle 竞赛中用作基准的训练/验证/测试划分,我们建议保留原始划分,以便将 Ludwig 模型与已发表的结果进行比较。

有关完整选项集,请参阅 ludwig.datasets.dataset_config.DatasetConfig。如果 DatasetConfig 提供的选项足以集成你的数据集,请跳到步骤 3:测试你的数据集。

但是,如果数据集需要默认数据集加载器未提供的其他处理,请继续执行步骤 2。

2. 如有需要,定义数据集加载器

如果 DatasetConfig 提供的选项不能涵盖你的数据集格式,或者数据集在训练前需要独特的处理,你可以在数据集加载器中添加 Python 代码。

加载器类应继承自 ludwig.datasets.loaders.dataset_loader.DatasetLoader,其模块名称应与数据集名称匹配。例如,AG News 在 ludwig/datasets/loaders/agnews.py 中有一个数据集加载器 agnews.AGNewsLoader

要指示 Ludwig 使用你的加载器,请将 loader 属性添加到你的数据集配置中

loader: agnews.AGNewsLoader

数据集的处理分为四个阶段

  1. 下载 - 数据集文件下载到缓存中。
  2. 验证 - 验证已下载文件的哈希值。
  3. 提取 - 从存档中提取数据集文件(如果数据未存档,则可能为空操作)。
  4. 转换 - 将数据集转换为可用于训练的格式,并准备加载。
    1. 转换文件(文件 -> 文件)
    2. 加载数据框(文件 -> 数据框)
    3. 转换数据框(数据框 -> 数据框)
    4. 保存处理后的数据(数据框 -> 文件)

对于每个阶段,在 ludwig.datasets.loaders.DatasetLoader 中都有一个相应的方法,可以重写该方法来提供自定义处理。

3. 测试你的数据集

创建一个简单的训练脚本和 ludwig 配置,以确保 Ludwig 训练 API 可与新数据集一起运行。例如

from ludwig.api import LudwigModel
from ludwig.datasets import titanic

training_set, test_set, _, = titanic.load(split=True)
model = LudwigModel(config="model_config.yaml", logging_level=logging.INFO)
train_stats, _, _ = model.train(training_set=training_set, test_set=test_set, model_name="titanic_model")

如果你添加了自定义加载器,请同时添加单元测试,以确保你的加载器在未来版本中仍然有效。按照以下示例,为单元测试提供少量数据样本,以便测试无需下载数据集。

单元测试示例


Kaggle 数据集的注意事项

为了测试从 Kaggle 下载数据集,请按照这些说明获取必要的 API 凭据。如果数据集是竞赛的一部分,你还需要在 Kaggle 网页界面中接受竞赛条款。

为了进行测试,Titanic 示例还演示了如何在测试中使用模拟的 kaggle 客户端。单元测试应该在没有凭据或互联网连接的情况下也能运行。


4. 添加建模示例

考虑分享一个关于用户如何使用你的数据集训练模型的示例,例如