Dataset类__getitem__报KeyError，如何正确映射字典型数据集？

直接解决方案 #

你的问题根源是DataLoader生成的是连续整数索引（0,1,2...），但你的数据集键是非连续的自定义整数（如939），直接用idx当键自然找不到对应样本。解决步骤如下：

1. 在数据集类的__init__方法中，把数据集的所有键提取成一个列表，建立连续索引到实际键的映射
2. 修改__getitem__，先通过idx获取对应的实际键，再从数据集中取样本

代码示例 #

from torch.utils.data import Dataset
from collections import defaultdict

class RefCOCOgDataset(Dataset):
    def __init__(self, dataset: defaultdict):
        self.dataset = dataset
        # 提取数据集的所有键，保存为列表用于索引映射
        self.sample_keys = list(self.dataset.keys())
    
    def __len__(self):
        # 返回总样本数，即键的数量
        return len(self.sample_keys)
    
    def __getitem__(self, idx: int):
        # 通过连续索引idx获取数据集的实际键
        current_key = self.sample_keys[idx]
        sample = self.dataset[current_key]
        
        # 这里可以根据需求处理样本（比如加载图片、格式化文本）
        image_filename = sample['image']
        texts = sample['texts']
        
        return {
            'image_filename': image_filename,
            'texts': texts
        }

__getitem__的迭代逻辑 #

DataLoader和Dataset的协作逻辑很直接：

• 首先调用__len__获取数据集总样本数，确定迭代的边界（比如总样本数1000，就会生成0到999的索引）
• 每次迭代时（按batch_size批量取样本），DataLoader会生成一组连续整数索引，逐个传给__getitem__
• __getitem__的唯一职责是：接收这个连续索引，返回对应位置的样本数据
• 你之前的错误就是跳过了"连续索引→实际键"的映射，直接把连续索引当成了数据集的键，而你的数据集键根本不是从0开始的连续整数，所以触发KeyError: 'image'（本质是找不到对应键的样本）

额外注意事项 #

• 确保你的defaultdict中所有键对应的样本都包含'image'字段，避免其他场景下的KeyError
• 如果数据集有新增/删除样本的情况，要同步更新self.sample_keys列表，否则会出现索引不匹配

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