HuggingFace：使用自定义 data_loader 和 data_collat​​or 从本地目录流式传输数据集

Posted 2023-03-29

【中文标题】HuggingFace：使用自定义 data_loader 和 data_collat​​or 从本地目录流式传输数据集

【英文标题】：HuggingFace: Streaming dataset from local dir using custom data_loader and data_collator

【发布时间】：2021-11-28 17:17:30

【问题描述】：

我有自定义 data_loader 和 data_collat​​or 用于使用 HuggingFace API 在 Transformer 模型中进行训练。它还进行数据集的映射，其中也进行了标记化。 我的 data_loader 脚本是一个固有的类 datasets.GeneratorBasedBuilder 所以包含 _generate_examples 产生样本的函数。

在开始训练时，它会缓存整个数据集（在系统上只缓存一次），然后开始训练。我可以在本地系统上重用该缓存，但不能在任何其他系统上使用该缓存的 .arrow 文件，因此缓存过程会重新启动。我想通过使用流媒体功能来避免缓存。我当前的代码如下：

from datasets import load_dataset
dataset = load_dataset ("/../my_data_loader.py", streaming =True)
train_dataset = dataset["train"]
train_datatset = train_dataset.map (..... )
data_collator = MyDataCollaor (......)
...
...
trainer = Trainer (model=model, arg= training_arg, train_dataset=train_dataset, data_collaor...)

注意：我不需要在我必须编码的地方：__len__ 和 __iter__ 函数在我身边。

【参考方案1】：

使用 1.12 或以上版本的数据集，我们可以通过设置streaming =True 来流式传输数据集（无缓存），如下所示。

 dataset = load_dataset ("/../my_data_loader.py", streaming =True)

在这种情况下，数据集将是 Iterable 数据集，因此 mapping 也将略有不同。假设以下脚本在caching 模式下使用：

train_dataset = datasets["train"]
train_dataset = train_dataset.map(
            tokenize_and_align_labels,
            batched=True,
            remove_columns=remove_columns,
            num_proc= preprocessing_num_workers,
            load_from_cache_file= not overwrite_cache,
        )

然后开启流式传输后，你就可以转换数据集格式，也可以改变映射函数的参数。

train_dataset = datasets["train"]
train_dataset = train_dataset.with_format("torch")
train_dataset = train_dataset.map(
        tokenize_and_align_labels,
        batched=True,
         )