PyTorch IMDb数据集仅返回负样本的问题排查求助

问题根源 #

IMDb数据集的test集默认存储顺序是先全部12500条负样本（标签1），再全部12500条正样本（标签2）。你的代码中先执行batch(2)再调用shuffle()，这会导致shuffle操作仅打乱整个batch的顺序，而非单个样本的顺序。前12500个batch全是负样本，你只取前10个batch，自然只会看到负样本。

修正后的代码 #

from torch.utils.data import DataLoader
from functools import partial
from torchtext.datasets import IMDB

imdb_datapipe = IMDB(split='test')

labels = {"1": "negative", "2": "positive"}
def process_labels(labels, x):
    return x[1], labels[str(x[0])]


# 调整操作顺序：先打乱单个样本，再按batch分组
imdb_datapipe = imdb_datapipe.map(partial(process_labels, labels))
imdb_datapipe = imdb_datapipe.shuffle()  # 先打乱所有样本的顺序
imdb_datapipe = imdb_datapipe.batch(2)   # 再按batch_size=2分组
imdb_datapipe = imdb_datapipe.rows2columnar(["text", "label"])
imdb_dataloader = DataLoader(imdb_datapipe, batch_size=None)

it = iter(imdb_dataloader)

for _ in range(10):
    sample = next(it)
    for text,label in zip(sample['text'], sample['label']):
        print(f"{label}: {text[:100]}")

补充说明 #

如果担心shuffle()的默认buffer_size不足以充分打乱数据，可以手动指定更大的buffer值，比如imdb_datapipe.shuffle(buffer_size=10000)，不过核心问题还是操作顺序错误，调整顺序后即可获取均衡的正负样本。

内容的提问来源于stack exchange，提问作者Tobias Strauß