PyTorch中使用DataLoader时如何访问下一迭代步骤的数据？

核心逻辑是手动控制DataLoader的迭代流程，提前缓存下一个步长的批次数据，不需要修改原有模型计算、数据加载的配置。

方法1：原生迭代器缓存（推荐，无额外依赖） #

这个方法完全兼容原有step计数规则，边界情况处理灵活，不会打乱DataLoader的采样、加载逻辑：

# 初始化DataLoader迭代器
dataloader_iter = iter(dataloader)
total_steps = len(dataloader)
# 加载第一个批次作为初始当前步数据
current_data = next(dataloader_iter)

for step in range(total_steps):
    # 预取下一批次，捕获迭代结束的异常
    try:
        next_data = next(dataloader_iter)
    except StopIteration:
        next_data = None  # 最后一步无后续批次，按业务需求自定义处理

    # 原有业务逻辑
    # ............................................................
    output = neuralnetwork_model(current_data)
    # 直接读取预取的下一步数据
    access = next_data
    # .............................................................

    # 更新当前批次为预取的下一批，进入下一轮迭代
    current_data = next_data

注意事项 #

• step计数和原enumerate(dataloader, 0)逻辑完全一致，从0开始逐次递增
• 最后一步（step = total_steps - 1）时next_data为None，可根据需求自定义处理逻辑，比如跳过访问、填充占位数据等
• DataLoader原有的shuffle、多进程加载（num_workers>0）、自定义采样器等配置全部正常生效，不会出现数据重复、跳步问题
• 仅额外占用一个批次的内存/显存，对训练资源的影响可以忽略

方法2：itertools工具实现（代码更简洁） #

如果不想手动写异常捕获，可以借助标准库的迭代器工具实现相邻批次偏移，需要注意最后一步的边界处理：

from itertools import tee, zip_longest

# 生成两个独立的DataLoader枚举迭代器
iter_current, iter_next = tee(enumerate(dataloader, 0))
# 第二个迭代器向前偏移一位
next(iter_next, None)

for (step, current_data), (_, next_data) in zip_longest(iter_current, iter_next):
    # 原有业务逻辑
    # ............................................................
    output = neuralnetwork_model(current_data)
    access = next_data  # 最后一步next_data为None
    # .............................................................

避坑提示 #

• 不要在原for循环内直接调用next(dataloader)，会打乱迭代器的指针位置，导致批次跳步、重复加载
• 不要尝试缓存整个DataLoader的所有批次到列表，会大幅提升内存占用，尤其是数据集较大时会直接导致内存溢出

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