项目简介
本项目是一个详细的教程,助力用户理解并上手PyTorch中的Distributed Data Parallel(DDP)技术。DDP作为PyTorch中用于多机多卡训练的关键工具,能大幅提升训练速度和模型规模。教程借助详尽的代码示例与理论阐释,协助用户从单机单卡训练逐步过渡到多机多卡的分布式训练。
项目的主要特性和功能
- 数据并行训练:详细讲解
nn.DataParallel
和nn.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】