如何延长PyTorch分布式训练的超时时间？

你的YOLOv5分布式训练在30分钟后中断，原因是PyTorch分布式初始化的默认超时时间为1800秒（30分钟）。以下两种方法可以解决这个问题：

方法一：通过环境变量快速调整 #

无需修改代码，在训练命令前添加TORCH_DIST_INIT_TIMEOUT环境变量，指定更长的超时时间（单位为秒）。例如设置为1小时（3600秒）：

TORCH_DIST_INIT_TIMEOUT=3600 python3 -m torch.distributed.run --nproc_per_node 2 train.py --batch 100 --epochs 1000 --data /home/username/Documents/folder_name/numbers.yaml --weights yolov5s.pt --device 0,1 --hyp data/hyps/hyp.scratch-high.yaml --name folder_name --patience 0

方法二：修改YOLOv5训练代码 #

如果需要长期固定配置，可修改YOLOv5中分布式初始化的代码逻辑。找到train.py或相关模块里调用torch.distributed.init_process_group的位置，添加timeout参数：

import torch.distributed as dist
from datetime import timedelta

# 保留原有初始化参数，新增timeout设置
dist.init_process_group(
    backend='nccl',  # 保持你当前使用的后端不变
    timeout=timedelta(seconds=3600)  # 设置为1小时，可按需调整时长
)

注意事项 #

• 超时时间需根据训练任务的实际耗时调整，确保能覆盖单轮训练、数据加载等环节的最长耗时。
• 环境变量方式适合临时调整；代码修改方式适合需要固定配置的场景。

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