分布式训练是将机器学习模型的训练过程分散到多个计算节点或设备上,以提高训练速度和效率,尤其是在处理大规模数据和模型时。分布式训练主要分为数据并行和模型并行两种主要策略:
数据并行是最常见的分布式训练方式。在这种方法中,模型副本会被复制到多个计算设备上,每个设备处理不同的批次(batch)数据。
# Replicate module to devices in device_ids
replicas = nn.parallel.replicate(module, device_ids)
# Distribute input to devices in device_ids
inputs = nn.parallel.scatter(input, device_ids)
# Apply the models to corresponding inputs
outputs = nn.parallel.parallel_apply(replicas, inputs)
# Gather result from all devices to output_device
result = nn.parallel.gather(outputs, output_device)
模型并行将一个大型模型拆分到多个设备上,以便更好地利用计算资源,尤其适用于内存消耗较大的模型。
tf.distribute.Strategy
轻松实现。torch.distributed
提供原生支持,还支持基于 Horovod 等第三方工具的分布式训练。常用的 API 有两个:
torch.nn.DataParallel
(简称 DP
)是 PyTorch 提供的一个简单的并行化工具,主要用于在多个 GPU 上进行数据并行训练。DataParallel
通过将输入数据批次(batch)切分成多个小批次,并将其分发到多个 GPU 上,进行并行处理。它会自动处理梯度的同步和模型参数的更新。
torch.nn.DataParallel
的工作机制DataParallel
会将模型复制到多个 GPU 上,每个 GPU 上有一个模型副本。cuda:0
)会收集所有 GPU 计算出的梯度,并将它们平均化,更新模型的全局参数。torch.nn.DataParallel
使用 DataParallel
非常简单,通常只需要将模型用 DataParallel
包裹,然后像普通模型一样使用即可。
import torch
import torch.nn as nn
import torch.optim as optim
# 定义模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
# 初始化模型和数据
model = SimpleModel()
# 将模型并行化
if torch.cuda.device_count() > 1:
print("Using", torch.cuda.device_count(), "GPUs")
model = nn.DataParallel(model)
model = model.cuda()
# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 模拟输入数据
inputs = torch.randn(32, 10).cuda() # 一个 32 样本的 batch,每个样本 10 个特征
targets = torch.randn(32, 5).cuda() # 对应的目标输出
# 前向传播
outputs = model(inputs)
# 计算损失
loss = criterion(outputs, targets)
# 反向传播
optimizer.zero_grad()
loss.backward()
# 更新模型参数
optimizer.step()
DistributedDataParallel
(简称 DDP
) 是 PyTorch 用于分布式训练的高级并行化工具,它的效率和灵活性比 DataParallel
更高,特别适合在多个 GPU 甚至跨多个节点(机器)上进行分布式训练。与 DataParallel
不同,DDP
在每个设备(GPU)上独立处理模型的前向传播和反向传播,并且避免了主设备的瓶颈问题。
DistributedDataParallel
的工作原理DataParallel
类似,DDP
会在每个 GPU 上保留一份模型副本。但与 DataParallel
不同的是,DDP
不需要将数据集中在主设备上,而是让每个 GPU 独立完成自己的工作。all-reduce
操作在所有设备之间同步,确保所有模型副本的梯度相同。这个过程是并行进行的,不会像 DataParallel
那样集中在主设备上,因此通信效率更高。DistributedDataParallel
的优点DDP
支持跨多台机器的训练,适合超大规模模型或需要跨节点的分布式训练。DDP
实现了无锁的梯度同步,不会因锁机制造成性能损失。DistributedDataParallel
的使用与 DataParallel
类似,DDP
也需要对模型进行包装,但它需要更多的设置,特别是在多机环境下,还需要配置通信后端。
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 初始化分布式环境
def setup(rank, world_size):
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
# 销毁分布式环境
def cleanup():
dist.destroy_process_group()
# 定义模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
# 初始化模型、优化器和数据
def main(rank, world_size):
setup(rank, world_size)
model = SimpleModel().cuda(rank)
ddp_model = DDP(model, device_ids=[rank])
criterion = nn.MSELoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)
# 模拟输入数据
inputs = torch.randn(32, 10).cuda(rank)
targets = torch.randn(32, 5).cuda(rank)
# 前向传播
outputs = ddp_model(inputs)
loss = criterion(outputs, targets)
# 反向传播
optimizer.zero_grad()
loss.backward()
# 更新模型参数
optimizer.step()
cleanup()
# 假设有两个GPU,可以这样启动分布式训练
if __name__ == "__main__":
world_size = 2 # GPU数
torch.multiprocessing.spawn(main, args=(world_size,), nprocs=world_size, join=True)
特性 | DataParallel (DP) | DistributedDataParallel (DDP) |
通信模式 | 主设备负责梯度同步 | 所有设备间并行同步梯度 |
性能 | 通信开销大,主设备瓶颈 | 通信开销小,性能更高 |
可扩展性 | 适用于单机多 GPU | 适用于单机或多机多 GPU |
使用场景 | 小规模并行 | 大规模或跨节点分布式训练 |
在深度学习任务中,数据加载通常是训练过程中的一个瓶颈,特别是当数据量很大时。使用多个进程来并行加载数据,并将数据从可分页内存(虚拟内存)转移到固定内存(GPU 内存)可以显著提高训练效率。
数据加载:
生产者-消费者模式:
固定内存的使用:
num_workers
:
num_workers
设置为大于 0 的值可以让 DataLoader
使用多个子进程来并行加载数据。num_workers=4
表示使用 4 个进程来加载数据。这可以显著提高数据加载速度,因为多个进程可以同时从磁盘读取不同的数据批次。pin_memory
:
pin_memory=True
时,DataLoader
会将数据从可分页的内存(虚拟内存)传输到固定内存中,这样在将数据转移到 GPU 时,数据传输速度会更快,因为固定内存可以避免页面交换的开销。这种方法结合了数据加载、数据传输和数据并行处理的优化,能够显著提升深度学习模型的训练效率和速度。
import torch
from torch.utils.data import DataLoader, Dataset
import numpy as np
class CustomDataset(Dataset):
def __init__(self, size):
self.data = np.random.rand(size, 3, 224, 224).astype(np.float32)
self.labels = np.random.randint(0, 2, size).astype(np.int64)
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return torch.tensor(self.data[idx]), torch.tensor(self.labels[idx])
dataset = CustomDataset(size=10000)
dataloader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=4, # 使用 4 个子进程加载数据
pin_memory=True # 将数据转移到固定内存
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
# 模型训练代码
# ...
参考文章:
Pytorch 分布式训练(DP/DDP)_pytorch分布式训练-CSDN博客https://blog.csdn.net/ytusdc/article/details/122091284?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522CC589E02-BBE1-4F15-BDC0-CA76EBF6C160%2522%252C%2522scm%2522%253A%252220140713.130102334..%2522%257D&request_id=CC589E02-BBE1-4F15-BDC0-CA76EBF6C160&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2~all~top_positive~default-1-122091284-null-null.142^v100^control&utm_term=%E5%88%86%E5%B8%83%E5%BC%8F%E8%AE%AD%E7%BB%83&spm=1018.2226.3001.4187