YOLO多机多GPU训练环境

原创已于 2025-09-28 10:28:46 修改 · 232 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #深度学习 #算法

于 2025-09-28 10:28:03 首次发布

部署运行你感兴趣的模型镜像

本文介绍如何使用多台电脑的GPU进行YOLO模型的分布式训练，适合初学者跟着操作。

准备工作

硬件要求

多台带GPU的电脑（GPU型号最好相同）
高速网络连接（建议千兆以上）
所有电脑在同一局域网下

软件要求

所有电脑需要安装相同版本的：

Python
PyTorch（带CUDA支持）
YOLO代码（如YOLOv5/YOLOv8）
NCCL库（用于多机通信）

环境配置步骤

1. 安装基础软件

在所有电脑上执行以下命令：

# 安装PyTorch
pip install torch torchvision torchaudio

# 下载YOLOv5
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

2. 配置SSH免密登录

选择一台作为主电脑，执行：

# 生成密钥（一直按回车即可）
ssh-keygen -t rsa

# 将密钥复制到所有电脑（包括自己）
ssh-copy-id username@电脑1的IP
ssh-copy-id username@电脑2的IP
...

启动分布式训练

在主电脑上运行：

python -m torch.distributed.launch \
    --nproc_per_node=4 \        # 每台电脑的GPU数量
    --nnodes=4 \                # 总电脑数量
    --node_rank=0 \             # 主电脑编号为0
    --master_addr="主电脑IP" \  # 主电脑的IP地址
    --master_port=12345 \       # 任意空闲端口号
    train.py \
    --batch-size 64 \           # 每GPU的批大小
    --data coco.yaml \          # 数据集配置
    --weights yolov5s.pt \      # 预训练权重
    --device 0,1,2,3 \          # 使用所有GPU
    --epochs 300

在其他电脑上运行：

命令与主电脑相同，只需将node_rank改为各自编号（1,2,3...）：

python -m torch.distributed.launch \
    --nproc_per_node=4 \
    --nnodes=4 \
    --node_rank=1 \            # 第二台电脑改为1，第三台改为2，以此类推
    --master_addr="主电脑IP" \
    --master_port=12345 \
    train.py \
    ...                        # 其他参数与主电脑保持一致