Pytorch使用教程(12)-如何进行并行训练？

在使用GPU训练大模型时，往往会面临单卡显存不足的情况。这时，通过多卡并行的形式来扩大显存是一个有效的解决方案。PyTorch主要提供了两个类来实现多卡并行：数据并行torch.nn.DataParallel（DP）和模型并行torch.nn.DistributedDataParallel（DDP）。本文将详细介绍这两种方法。

一、数据并行（torch.nn.DataParallel）

1. 基本原理
   数据并行是一种简单的多GPU并行训练方式。它通过多线程的方式，将输入数据分割成多个部分，每个部分在不同的GPU上并行处理，最后将所有GPU的输出结果汇总，计算损失和梯度，更新模型参数。
   

2. 使用方法
   使用torch.nn.DataParallel非常简单，只需要一行代码就可以实现。以下是一个示例：

import torch
import torch.nn as nn# 检查是否有多个GPU可用
if torch.cuda.device_count() > 1:print("Let's use", torch.cuda.device_count(), "GPUs!")# 将模型转换为DataParallel对象model = nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))

3. 优缺点
   ‌优点‌：代码简单，易于使用，对小白比较友好。
   ‌缺点‌：GPU会出现负载不均衡的问题，一个GPU可能占用了大部分负载，而其他GPU却负载较轻，导致显存使用不平衡。

二、模型并行（torch.nn.DistributedDataParallel）

1. 基本原理
   torch.nn.DistributedDataParallel（DDP）是一种真正的多进程并行训练方式。每个进程对应一个独立的训练过程，且只对梯度等少量数据进行信息交换。每个进程包含独立的解释器和GIL（全局解释器锁），因此可以充分利用多GPU的优势，实现更高效的并行训练。
   

2. 使用方法

   使用torch.nn.DistributedDataParallel需要进行一些额外的配置，包括初始化GPU通信方式、设置随机种子点、使用DistributedSampler分配数据等。以下是一个详细的示例：

初始化环境

import torch
import torch.distributed as dist
import argparsedef parse():parser = argparse.ArgumentParser()parser.add_argument('--local_rank', type=int, default=0)args = parser.parse_args()return argsdef main():args = parse()torch.cuda.set_device(args.local_rank)dist.init_process_group('nccl', init_method='env://')device = torch.device(f'cuda:{args.local_rank}')

设置随机种子点

import numpy as np# 固定随机种子点
seed = np.random.randint(1, 10000)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)使用DistributedSampler分配数据
python
Copy Code
from torch.utils.data.distributed import DistributedSamplertrain_dataset = ...  # 你的数据集
train_sampler = DistributedSampler(train_dataset, shuffle=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opts.batch_size, sampler=train_sampler
)

初始化模型

model = mymodel().to(device)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])训练循环
python
Copy Code
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()for ep in range(total_epoch):train_sampler.set_epoch(ep)for inputs, labels in train_loader:inputs, labels = inputs.to(device), labels.to(device)optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()

3. 优缺点

• 优点‌：每个进程对应一个独立的训练过程，显存使用更均衡，性能更优。
• 缺点‌：代码相对复杂，需要进行一些额外的配置。

三、对比与选择

1. 对比

特点	torch.nn.DataParallel	torch.nn.DistributedDataParallel
并行方式	多线程	多进程
显存使用	可能不均衡	更均衡
性能	一般	更优
代码复杂度	简单	复杂

3. 选择建议

• 对于初学者或快速实验，可以选择torch.nn.DataParallel，因为它代码简单，易于使用。
• 对于需要高效并行训练的场景，建议选择torch.nn.DistributedDataParallel，因为它可以充分利用多GPU的优势，实现更高效的训练。

四、小结

通过本文的介绍，相信读者已经对PyTorch的多GPU并行训练有了更深入的了解。在实际应用中，可以根据模型的复杂性和数据的大小选择合适的并行训练方式，并调整batch size和学习率等参数以优化模型的性能。希望这篇文章能帮助你掌握PyTorch的多GPU并行训练技术。