Pytorch使用教程(12)-如何进行并行训练？
在使用GPU训练大模型时，往往会面临单卡显存不足的情况。这时，通过多卡并行的形式来扩大显存是一个有效的解决方案。PyTorch主要提供了两个类来实现多卡并行：数据并行torch.nn.DataParallel（DP）和模型并行torch.nn.DistributedDataParallel（DDP）。本文将详细介绍这两种方法。
一、数据并行（torch.nn.DataParallel）
-
基本原理
数据并行是一种简单的多GPU并行训练方式。它通过多线程的方式，将输入数据分割成多个部分，每个部分在不同的GPU上并行处理，最后将所有GPU的输出结果汇总，计算损失和梯度，更新模型参数。
-
使用方法
使用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()))
- 优缺点
优点：代码简单，易于使用，对小白比较友好。
缺点：GPU会出现负载不均衡的问题，一个GPU可能占用了大部分负载，而其他GPU却负载较轻，导致显存使用不平衡。
二、模型并行（torch.nn.DistributedDataParallel）
-
基本原理
torch.nn.DistributedDataParallel（DDP）是一种真正的多进程并行训练方式。每个进程对应一个独立的训练过程，且只对梯度等少量数据进行信息交换。每个进程包含独立的解释器和GIL（全局解释器锁），因此可以充分利用多GPU的优势，实现更高效的并行训练。
-
使用方法
使用torch.nn.DistributedDataParallel需要进行一些额外的配置，包括初始化GPU通信方式、设置随机种子点、使用DistributedSampler分配数据等。以下是一个详细的示例：
初始化环境
import torch
import torch.distributed as dist
import argparse
def parse():
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
return args
def 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 DistributedSampler
train_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()
- 优缺点
- 优点：每个进程对应一个独立的训练过程，显存使用更均衡，性能更优。
- 缺点：代码相对复杂，需要进行一些额外的配置。
三、对比与选择
- 对比
| 特点 | torch.nn.DataParallel | torch.nn.DistributedDataParallel |
|---|---|---|
| 并行方式 | 多线程 | 多进程 |
| 显存使用 | 可能不均衡 | 更均衡 |
| 性能 | 一般 | 更优 |
| 代码复杂度 | 简单 | 复杂 |
- 选择建议
- 对于初学者或快速实验，可以选择torch.nn.DataParallel，因为它代码简单，易于使用。
- 对于需要高效并行训练的场景，建议选择torch.nn.DistributedDataParallel，因为它可以充分利用多GPU的优势，实现更高效的训练。
四、小结
通过本文的介绍，相信读者已经对PyTorch的多GPU并行训练有了更深入的了解。在实际应用中，可以根据模型的复杂性和数据的大小选择合适的并行训练方式，并调整batch size和学习率等参数以优化模型的性能。希望这篇文章能帮助你掌握PyTorch的多GPU并行训练技术。
原文地址：https://blog.csdn.net/evm_doc/article/details/145244408
免责声明：本站文章内容转载自网络资源，如侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！