HuggingFace Datasets库中数据集保存与加载的正确方式

在使用HuggingFace的Datasets库处理数据时，许多开发者会遇到一个常见问题：保存的数据集在重新加载后大小不一致。本文将深入分析这一现象的原因，并详细介绍正确的数据集保存与加载方法。

问题现象

当开发者使用save_to_disk方法保存数据集后，如果错误地使用load_dataset而不是load_from_disk来加载数据，会导致加载后的数据集大小与原始数据集不符。例如：

• 保存时：训练集7642条，测试集1000条
• 错误加载后：训练集和测试集都只有1条

原因分析

这种现象源于Datasets库设计的两种不同数据加载机制：

1. load_from_disk：专门用于加载通过save_to_disk保存的二进制格式数据集
2. load_dataset：用于从原始数据文件(如JSON、CSV等)创建新的数据集

当使用load_dataset加载二进制格式的数据目录时，库会尝试将其解释为原始数据文件目录，导致数据解析错误。

正确的保存与加载方式

方法一：使用专用二进制格式

# 保存数据集
dataset.save_to_disk("dataset_path")

# 加载数据集
from datasets import load_from_disk
dataset = load_from_disk("dataset_path")

这种方法保存的是经过优化的二进制格式，具有以下优势：

• 加载速度快
• 保留所有数据集特征和元数据
• 支持内存映射，节省内存

方法二：使用原始文件格式

如果希望数据集目录可以直接用load_dataset加载，需要按照特定结构组织原始数据文件：

dataset_dir/
├── train.jsonl
└── test.jsonl

保存代码示例：

import os
import shutil

os.makedirs("dataset_dir", exist_ok=True)
shutil.copy("train_data.jsonl", "dataset_dir/train.jsonl")
shutil.copy("test_data.jsonl", "dataset_dir/test.jsonl")

加载代码示例：

from datasets import load_dataset
dataset = load_dataset("dataset_dir")

两种方法的比较

特性	二进制格式(save_to_disk)	原始文件格式
加载速度	快	慢
存储空间	较小	较大
可读性	二进制不可读	原始格式可读
兼容性	需特定方法加载	通用格式
元数据保留	完整保留	可能丢失

最佳实践建议

1. 在训练流程中建议使用二进制格式，提高效率
2. 当需要与其他工具共享数据时，可使用原始文件格式
3. 大型数据集优先考虑二进制格式的内存映射优势
4. 调试阶段可以使用原始文件格式便于检查数据

通过理解这些差异并选择适当的方法，开发者可以避免数据集大小不一致的问题，并优化数据处理流程的效率。