数据集 🤝 Arrow
什么是 Arrow?
Arrow 使得能够快速处理和移动大量数据。它是一种特定的数据格式,以列式内存布局存储数据。这提供了几个显著的优势
- Arrow 的标准格式允许 零拷贝读取,从而几乎消除了所有序列化开销。
- Arrow 与语言无关,因此它支持不同的编程语言。
- Arrow 是面向列的,因此在查询和处理数据切片或列方面速度更快。
- Arrow 允许无复制地传递到标准机器学习工具,例如 NumPy、Pandas、PyTorch 和 TensorFlow。
- Arrow 支持许多(可能嵌套的)列类型。
内存映射
🤗 Datasets 使用 Arrow 作为其本地缓存系统。它允许数据集以磁盘缓存为后盾,该缓存被内存映射以实现快速查找。这种架构允许在设备内存相对较小的机器上使用大型数据集。
例如,加载完整的英文维基百科数据集只需要几 MB 的 RAM。
>>> import os; import psutil; import timeit
>>> from datasets import load_dataset
# Process.memory_info is expressed in bytes, so convert to megabytes
>>> mem_before = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
>>> wiki = load_dataset("wikipedia", "20220301.en", split="train")
>>> mem_after = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
>>> print(f"RAM memory used: {(mem_after - mem_before)} MB")
RAM memory used: 50 MB
这是可能的,因为 Arrow 数据实际上是从磁盘内存映射的,而不是加载到内存中。内存映射允许访问磁盘上的数据,并利用虚拟内存功能实现快速查找。
性能
使用 Arrow 迭代内存映射数据集速度很快。在笔记本电脑上迭代维基百科可以达到 1-3 Gbit/s 的速度。
>>> s = """batch_size = 1000
... for batch in wiki.iter(batch_size):
... ...
... """
>>> elapsed_time = timeit.timeit(stmt=s, number=1, globals=globals())
>>> print(f"Time to iterate over the {wiki.dataset_size >> 30} GB dataset: {elapsed_time:.1f} sec, "
... f"ie. {float(wiki.dataset_size >> 27)/elapsed_time:.1f} Gb/s")
Time to iterate over the 18 GB dataset: 31.8 sec, ie. 4.8 Gb/s