在Linux环境下使用PyTorch进行分布式训练,可以遵循以下步骤:
安装PyTorch: 确保你已经安装了PyTorch,并且版本支持分布式训练。
设置环境变量:
为了启用分布式训练,需要设置一些环境变量,例如NCCL_DEBUG=INFO
和HOROVOD_TIMELINE
。
网络配置: 确保所有参与分布式训练的节点可以通过SSH无密码登录,并且网络配置正确。
使用torch.distributed.launch
:
PyTorch提供了一个脚本torch.distributed.launch
来启动分布式训练。这个脚本会自动处理一些初始化工作。
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py
其中,NUM_GPUS_YOU_HAVE
是你每台机器上的GPU数量,YOUR_TRAINING_SCRIPT.py
是你的训练脚本。
使用horovodrun
:
如果你使用Horovod进行分布式训练,可以使用horovodrun
命令来启动训练。
horovodrun -np NUM_CPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py
其中,NUM_CPUS_YOU_HAVE
是你每台机器上的CPU数量。
在你的训练脚本中,需要进行以下修改:
初始化分布式环境:
使用torch.distributed.init_process_group
来初始化分布式环境。
import torch.distributed as dist
dist.init_process_group(
backend='nccl', # 或者 'gloo'
init_method='tcp://<master_ip>:<master_port>',
world_size=<world_size>,
rank=<rank>
)
其中,<master_ip>
是主节点的IP地址,<master_port>
是主节点的端口号,<world_size>
是总的进程数,<rank>
是当前进程的排名。
数据并行:
使用torch.nn.parallel.DistributedDataParallel
来包装你的模型。
model = YourModel().to(rank)
ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])
数据加载器:
确保你的数据加载器支持分布式训练。可以使用torch.utils.data.distributed.DistributedSampler
。
from torch.utils.data import DataLoader, DistributedSampler
train_dataset = YourDataset()
train_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, batch_size=your_batch_size, sampler=train_sampler)
以下是一个简单的分布式训练示例:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, DistributedSampler
import torchvision.datasets as datasets
import torchvision.transforms as transforms
# 初始化分布式环境
dist.init_process_group(
backend='nccl',
init_method='tcp://<master_ip>:<master_port>',
world_size=<world_size>,
rank=<rank>
)
# 定义模型
class YourModel(nn.Module):
def __init__(self):
super(YourModel, self).__init__()
# 定义你的模型结构
def forward(self, x):
# 定义前向传播
return x
model = YourModel().to(rank)
ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)
# 加载数据
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, batch_size=64, sampler=train_sampler)
# 训练循环
for epoch in range(num_epochs):
train_sampler.set_epoch(epoch)
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
inputs, labels = data[0].to(rank), data[1].to(rank)
optimizer.zero_grad()
outputs = ddp_model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')
# 清理分布式环境
dist.destroy_process_group()
通过以上步骤,你可以在Linux环境下使用PyTorch进行分布式训练。确保所有节点的网络配置正确,并且环境变量设置无误。