运行train.py训练神经网络时出现StopIteration错误的解决求助

问题根因 #

• 启动参数中num_gpus=0，你在修改代码时直接将DataLoader的batch_size设为args.num_gpus，导致batch size为0，DataLoader无有效数据可以加载，迭代器为空直接抛出StopIteration
• 迭代器iterator_train是在epoch循环外层初始化的，第一个epoch跑完数据集后迭代器就已经耗尽，第二个及之后的epoch调用train函数时直接读取已耗尽的迭代器就会报错
• 40个worker数远超CPU训练的合理负载，进程间通信异常触发段错误

修复方案 #

1. 修正DataLoader配置 #

替换你修改后的DataLoader代码为以下内容：

# 计算合法batch size，CPU训练时按单卡batch size计算
train_batch_size = args.batch_size_per_gpu * max(args.num_gpus, 1)
# CPU训练worker数建议设为2-4，最高不超过8，避免进程溢出
train_worker_num = min(int(args.workers), 4) if args.num_gpus == 0 else min(int(args.workers), 8)

loader_train = torchdata.DataLoader(
    dataset_train,
    batch_size=train_batch_size,
    shuffle=False,
    collate_fn=user_scattered_collate,
    num_workers=train_worker_num,
    drop_last=True,
    pin_memory=args.num_gpus > 0  # CPU训练不需要开启pin_memory，避免多余内存占用
)

2. 调整迭代器初始化逻辑 #

删除原来从外部传入iterator_train的逻辑，直接在每个epoch的train函数开头初始化迭代器，修改train函数开头代码：

def train(segmentation_module, optimizers, history, epoch, args, dataset_train):
    # 每个epoch开始前重新生成迭代器，避免上一轮耗尽的问题
    iterator = iter(loader_train)
    
    tic = time.time()
    for i in range(args.epoch_iters):
        batch_data = next(iterator)
        # 后面接原有训练逻辑

同时修改外层epoch循环调用train函数的代码，去掉iterator_train传参：

for epoch in range(args.start_epoch, args.num_epoch + 1):
    train(segmentation_module, optimizers, history, epoch, args, dataset_train)
    # 后面接原有 checkpoint、评估逻辑

3. 可选：兼容迭代器耗尽场景 #

如果你的epoch_iters设置大于数据集实际可迭代的batch数，可以保留try-except逻辑，调整为：

for i in range(args.epoch_iters):
    try:
        batch_data = next(iterator)
    except StopIteration:
        # 迭代器耗尽后重新生成
        iterator = iter(loader_train)
        batch_data = next(iterator)
    # 后续逻辑不变

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