Python多进程双GPU神经网络训练无报错异常退出问题求助

异常原因分析 #

核心原因是运行训练任务0的worker进程在执行split2的predict步骤后异常崩溃退出，触发父进程manage_gpus的worker.join()直接返回，于是父进程从队列取下一个任务2执行，所以没有打印split2结束和任务0完成的日志，并非代码逻辑本身有问题。

触发崩溃的常见原因如下：

• 显存溢出(OOM)：你虽然删除了无用变量，但PyTorch/TensorFlow等深度学习框架默认会持有独立的显存缓存池，删除Python对象不会立刻将显存释放回GPU设备，多轮训练累计的显存碎片或者峰值占用超过GPU显存上限后，进程会被系统直接杀死。
• 环境变量设置错误：你在worker进程内部才设置CUDA_VISIBLE_DEVICES，如果主进程提前导入了深度学习框架，框架会提前读取环境变量，导致两个worker进程实际都占用默认的GPU0，造成显存抢占，GPU0先触发OOM崩溃。
• 静默异常：训练代码的predict步骤或者后续清理逻辑触发了未捕获的异常，进程直接退出，没有打印错误信息。

修复方案 #

1. 新增异常捕获：在train_and_export_wrapper函数最外层加try-except块，捕获所有异常并打印堆栈，确认崩溃的具体原因：

def train_and_export_wrapper(trainingUid, device, lock):
    try:
        # 原有逻辑全部放在此处
    except Exception as e:
        import traceback
        with lock:
            print(f"Device {device} training {trainingUid} error: {str(e)}")
            print(traceback.format_exc())

2. 调整CUDA环境变量设置时机：将os.environ["CUDA_VISIBLE_DEVICES"]的设置移到worker进程最开头，且保证主进程在创建子进程前不要导入任何深度学习框架，或者在创建Process时通过env参数直接传入环境变量。
3. 强制释放显存：在每折训练删除变量后，调用框架的显存释放接口，比如PyTorch添加torch.cuda.empty_cache()，TensorFlow添加tf.keras.backend.clear_session()。

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