PyTorch自定义Dataset直接遍历H5存储BERT嵌入时报索引越界错误

核心逻辑认知纠正 #

torch.utils.data.Dataset是Map风格数据集的抽象基类，本身完全没有内置实现__iter__迭代方法。直接通过for循环遍历Dataset实例时，走的是Python原生的迭代协议，整个流程完全不读取__len__的返回值，不会按照__len__定义的边界控制遍历范围。
Python原生对实现了整数索引__getitem__方法的类的迭代规则为：从idx=0开始，每次索引自增1，调用__getitem__(idx)获取返回值，只要不抛出IndexError就持续遍历，直到捕获到IndexError才终止循环。这是所有同类异常的核心来源。

不同存储格式表现差异的原因 #

• 之前使用CSV存储的数据集没有出现越界问题，属于场景巧合：CSV场景下__getitem__内部通常直接读取内存中长度和样本数完全匹配的Python列表，当迭代器走到idx=列表长度的位置时，列表自身会抛出IndexError，刚好在样本边界终止遍历，容易让人误以为遍历流程是受__len__控制的。
• H5存储场景下，__getitem__内部直接调用h5py.Dataset的索引接口，这个对象的有效索引范围是0 ~ 22424（对应shape为(22425, 128, 768)）。无论你手动把__len__的返回值调整为100、22424还是其他数值，迭代器都会从0开始持续递增索引，直到idx=22425时触发h5py的越界报错——这也是调整__len__完全不影响报错触发点的根本原因。

可通过最小案例复现验证该逻辑：构造一个自定义类，实现__getitem__读取长度为1000的列表，同时故意将__len__返回值设为100，直接遍历这个类实例时，依然会完整遍历全部1000个元素，与__len__的返回值无关联。

DataLoader运行正常的原因 #

PyTorch内置的DataLoader对Map风格Dataset的处理逻辑和直接遍历完全独立：它根本不会调用Dataset实例的原生迭代方法，而是先主动调用__len__拿到总样本数N，自行生成0 ~ N-1的完整索引序列，逐个将索引传给__getitem__获取样本。整个流程完全受__len__的返回值约束，自然不会触碰到超出范围的索引。

修复方案 #

给自定义Dataset类手动实现__iter__方法，按照__len__的返回值生成合法索引序列，替代Python自动生成的不受控迭代逻辑，示例代码：

def __iter__(self):
    # 严格按照__len__定义的边界生成迭代序列
    for idx in range(len(self)):
        yield self[idx]

附加优化：h5py打开的文件对象不会随Dataset实例的销毁自动关闭，建议补充文件句柄的生命周期管理逻辑，避免长期运行时出现文件描述符泄漏。

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