PyTorch中MaskRCNN训练GPU环境下损失无法复现问题求助

针对你在Kaggle Notebook中使用maskrcnn_resnet50_fpn_v2时遇到的GPU训练不可复现问题，结合单类大实例数据集的特点，可按以下步骤逐一排查解决：

• 消除前向传播的隐性随机性
  MaskRCNN的部分组件（如ROI池化、NMS采样）在GPU上默认采用非确定性实现，即使模型权重不变，两次前向传播也会产生差异。可手动固定相关参数：

  # 固定RPN和ROI Heads的NMS阈值，避免动态随机调整
  model.rpn.nms_thresh = 0.7
  model.roi_heads.nms_thresh = 0.7
  # 强制ROI池化使用确定性插值方式
  if hasattr(model.roi_heads.box_roi_pool, 'mode'):
      model.roi_heads.box_roi_pool.mode = 'bilinear'
  

  同时检查数据预处理环节：确保所有随机增强（如翻转、裁剪）完全关闭，单图测试时用固定的预处理流程，避免隐性随机操作。

• 精细化设置GPU环境种子
  仅用seed_everything无法覆盖PyTorch GPU组件的所有随机源，需补充以下设置（比torch.use_deterministic_algorithms(True)显存占用低很多）：

  import torch
  import numpy as np
  import random
  import os
  
  def set_seeds(seed=42):
      random.seed(seed)
      np.random.seed(seed)
      torch.manual_seed(seed)
      torch.cuda.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
      # 禁用cuDNN非确定性算法，关闭自动卷积优化
      torch.backends.cudnn.deterministic = True
      torch.backends.cudnn.benchmark = False
  

  每次启动训练前都要调用该函数，确保所有随机源完全固定。

• 解决第二epoch开始的损失差异
  第一epoch一致、后续epoch差异的核心原因是训练流程中的动态状态残留或随机采样：

  • 每次训练前重新初始化模型和优化器，避免上一轮训练的状态残留：

    # 重新创建模型与优化器，不复用旧实例
    model = maskrcnn_resnet50_fpn_v2(num_classes=2)  # 单类+背景，num_classes设为2
    model.to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)
    set_seeds(42)
    

  • 修改DataLoader采样器与工作进程：用SequentialSampler替代随机采样，同时将num_workers设为0，避免多进程数据加载的随机性：

    from torch.utils.data import DataLoader, SequentialSampler
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=1,
        sampler=SequentialSampler(train_dataset),
        num_workers=0,
        collate_fn=lambda x: tuple(zip(*x))
    )
    

• 适配单类大实例的采样设置
  单图1000个实例会导致默认ROI采样的随机性被放大，需固定采样参数：

  # 固定每个图像的ROI批次大小与正样本比例
  model.roi_heads.batch_size_per_image = 256
  model.roi_heads.positive_fraction = 0.25
  # 关闭动态随机采样，改用固定逻辑
  model.roi_heads.sampler = None
  

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