littlebot
Published on 2025-04-11 / 0 Visits
0

【源码】基于PyTorch框架的分布式数据并行训练教程

项目简介

本项目是一个详细的教程,助力用户理解并上手PyTorch中的Distributed Data Parallel(DDP)技术。DDP作为PyTorch中用于多机多卡训练的关键工具,能大幅提升训练速度和模型规模。教程借助详尽的代码示例与理论阐释,协助用户从单机单卡训练逐步过渡到多机多卡的分布式训练。

项目的主要特性和功能

  • 数据并行训练:详细讲解nn.DataParallelnn.parallel.DistributedDataParallel的区别与使用场景。
  • 多进程并行:通过DDP实现多进程并行训练,支持单机多卡和多机多卡训练模式。
  • 混合精度训练:介绍在DDP中启用混合精度训练以加速训练过程的方法。
  • 模型验证:提供在分布式环境中进行模型验证的示例代码。
  • 通信协议选择:介绍依据网络环境和需求选择合适通信协议(如NCCL和Gloo)的方法。

安装使用步骤

1. 环境准备

确保已安装PyTorch和torchvision库,并配置好CUDA环境(若使用GPU)。 bash pip install torch torchvision

2. 下载项目源码

用户已完成此步骤。

3. 单机单卡训练

运行以下命令启动单机单卡训练: bash python train.py -g 0

4. 多机多卡训练

使用TCP模式启动

在每台机器上分别启动进程,指定不同的rank和GPU ID: ```bash python mnist-tcp.py --init_method tcp://192.168.1.201:12345 -g 0 --rank 0 --world_size 4 --use_mix_precision python mnist-tcp.py --init_method tcp://192.168.1.201:12345 -g 1 --rank 1 --world_size 4 --use_mix_precision

python tcp_init.py --init_method tcp://192.168.1.201:12345 -g 0 --rank 2 --world_size 4 --use_mix_precision python tcp_init.py --init_method tcp://192.168.1.201:12345 -g 1 --rank 3 --world_size 4 --use_mix_precision ```

使用ENV模式启动

使用torch.distributed.launch脚本启动多机多卡训练: ```bash python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr="192.168.1.201" --master_port=12345 mnist-env.py --use_mix_precision

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=1 --master_addr="192.168.1.201" --master_port=12345 mnist-env.py --use_mix_precision ```

5. 模型验证

在训练过程中,可通过evaluate函数在验证集上进行模型验证。验证代码会自动在分布式环境中汇总验证结果。 python def evaluate(model, gpu, test_loader, rank): model.eval() size = torch.tensor(0.).to(gpu) correct = torch.tensor(0.).to(gpu) with torch.no_grad(): for i, (images, labels) in enumerate(tqdm(test_loader)): images = images.to(gpu) labels = labels.to(gpu) outputs = model(images) size += images.shape[0] correct += (outputs.argmax(1) == labels).type(torch.float).sum() dist.reduce(size, 0, op=dist.ReduceOp.SUM) dist.reduce(correct, 0, op=dist.ReduceOp.SUM) if rank==0: print('Evaluate accuracy is {:.2f}'.format(correct / size)) 通过以上步骤,可顺利进行基于PyTorch的分布式数据并行训练,并验证模型的性能。

下载地址

点击下载 【提取码: 4003】【解压密码: www.makuang.net】