Google Colab半监督EfficientNet迁移训练RAM逐轮上涨问题求助

核心泄漏点 #

• 旧版本的concat_dataset、train_loader、new_dataset未被手动释放：每轮半监督逻辑都会新建这三个实例，上一轮的实例如果没有被显式删除，会持续持有伪标签数据的内存引用，无法被GC回收
• 伪标签推理阶段的中间张量未释放：logits、probs、每个batch的img张量在推理完成后没有被删除，即使使用了torch.no_grad()，残留的张量对象仍会占用内存
• numpy数组与PyTorch张量共享内存：直接调用img[i].numpy()生成的数组和原PyTorch张量共享内存，只要numpy数组还被持有，原张量的内存就无法被释放
• 半监督逻辑结尾仅删除了dataset_0、dataset_1等中间变量，没有删除持有最终伪标签数据的new_dataset，导致对应的内存始终被占用

修复步骤 #

1. 添加上一轮旧对象的释放逻辑
   在每轮训练、验证流程完成后，下一轮epoch开始前，显式删除上一轮的相关实例：

# 放在训练循环、验证循环结束后，下一轮epoch开始前
del train_loader, concat_dataset, new_dataset
gc.collect()
# 若使用GPU，加一行清空CUDA缓存
torch.cuda.empty_cache()

2. 清理伪标签推理阶段的中间变量
   修改伪标签推理的循环，每个batch处理完成后立即删除中间张量，同时新增copy()断开numpy数组与PyTorch张量的内存共享：

for img in pseudo_loader:
    with torch.no_grad():
        logits = model(img.to(device))
    probs = softmax(logits)

    # 原有过滤逻辑不变
    for i in range(len(probs)):
        p = probs[i].tolist()
        idx = p.index(max(p))
        if p[idx] >= threshold:
            if idx == 0:
                dataset_0.append(img[i].numpy().copy().reshape(128, 128, 3))
            else:
                dataset_1.append(img[i].numpy().copy().reshape(128, 128, 3))
    # 每个batch处理完立即清理中间变量
    del logits, probs, img
    gc.collect()

3. 完善半监督模块结尾的释放逻辑
   在删除dataset_0等变量时，补充对不需要的临时对象的清理，确保所有不再使用的大内存对象都被显式标记为可回收。
4. 排查DataLoader配置
   如果pseudo_loader的num_workers参数大于0，可以尝试设置为num_workers=0，多进程加载数据在Colab环境下偶现内存泄漏的问题。

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