在HuggingFace Pipelines中为何使用KeyDataset而非直接调用dataset？

先看两种写法：

用KeyDataset的写法 #

from transformers import pipeline
from transformers.pipelines.pt_utils import KeyDataset
import datasets

dataset = datasets.load_dataset("imdb", name="plain_text", split="unsupervised")
pipe = pipeline("text-classification", device=0)
for out in pipe(KeyDataset(dataset, "text"), batch_size=8, truncation="only_first"):
    print(out)

直接调用dataset["text"]的写法 #

from transformers import pipeline
import datasets

dataset = datasets.load_dataset("imdb", name="plain_text", split="unsupervised")
pipe = pipeline("text-classification", device=0)
for out in pipe(dataset["text"], batch_size=8, truncation="only_first"):
    print(out)

核心区别在于这几点：

• 内存占用天差地别：dataset["text"]会把整个文本列一次性加载到内存里，要是数据集有几十万甚至上百万条数据，直接就会吃掉几GB内存；而KeyDataset是惰性加载的迭代器，只在处理当前批次的时候才去取对应的数据，全程几乎不占额外内存，处理超大数据集时优势特别明显。
• 流式处理更高效：KeyDataset实现了Hugging Face Dataset的迭代接口，配合pipeline的批量处理可以做到边取数据边处理，不用等所有数据加载完再开始计算，尤其在GPU上跑的时候，能更流畅地利用显存，避免因为内存溢出中断程序。
• 格式适配更友好：pipeline对迭代器类型的输入支持更原生，直接传dataset["text"]（本质是列表类对象）时，pipeline会先把所有数据预加载再分批次，而用KeyDataset可以跳过预加载步骤，直接按批次迭代获取数据，减少不必要的中间开销。

内容的提问来源于stack exchange，提问作者PS1