在CentOS上实现PyTorch的并行计算，主要依赖于PyTorch提供的多GPU支持。以下是详细的步骤和示例代码，帮助你在CentOS上配置和使用PyTorch进行并行计算。
首先，确保你已经安装了PyTorch。你可以使用pip或conda来安装PyTorch。以下是使用pip安装的示例：
pip install torch torchvision torchaudio
如果你需要GPU支持，请确保你的系统已经安装了CUDA和cuDNN，并选择相应的PyTorch版本。例如，如果你使用的是CUDA 11.7，可以这样安装：
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
PyTorch提供了多种方式来实现多GPU并行计算，包括数据并行（DataParallel）和分布式并行（DistributedParallel）。以下是这两种方法的简要介绍和示例代码。
数据并行是将输入数据分割到多个GPU上进行处理，然后将结果合并。适用于模型较小且数据量较大的情况。
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
# 定义一个简单的卷积神经网络
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
self.relu = nn.ReLU()
self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.fc1 = nn.Linear(32 * 14 * 14, 10)
def forward(self, x):
x = self.pool(self.relu(self.conv1(x)))
x = x.view(-1, 32 * 14 * 14)
x = self.fc1(x)
return x
# 创建模型实例
model = SimpleCNN()
# 使用DataParallel包装模型
if torch.cuda.device_count() > 1:
print(f"Let's use {torch.cuda.device_count()} GPUs!")
model = nn.DataParallel(model)
# 将模型移动到GPU
model.to('cuda')
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 加载数据集
transform = transforms.Compose([transforms.ToTensor()])
trainset = datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
# 训练模型
for epoch in range(5): # 多次循环遍历数据集
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data[0].to('cuda'), data[1].to('cuda')
# 梯度清零
optimizer.zero_grad()
# 前向传播
outputs = model(inputs)
# 计算损失
loss = criterion(outputs, labels)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
running_loss += loss.item()
if i % 100 == 99: # 每100个mini-batches打印一次
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 100:.3f}')
running_loss = 0.0
print('Finished Training')
分布式并行适用于大规模数据和模型，可以更高效地利用多台机器或多GPU系统。
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torchvision import datasets, transforms
import torch.distributed as dist
import torch.multiprocessing as mp
def train(rank, world_size):
# 初始化进程组
dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)
# 创建模型实例
model = SimpleCNN().to(rank)
# 使用DistributedDataParallel包装模型
ddp_model = DDP(model, device_ids=[rank])
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)
# 加载数据集
transform = transforms.Compose([transforms.ToTensor()])
trainset = datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=True, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset, num_replicas=world_size, rank=rank)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, sampler=train_sampler)
# 训练模型
for epoch in range(5): # 多次循环遍历数据集
train_sampler.set_epoch(epoch)
running_loss = 0.0
for i, data in enumerate(trainloader, 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()
if i % 100 == 99: # 每100个mini-batches打印一次
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 100:.3f}')
running_loss = 0.0
dist.destroy_process_group()
def main():
world_size = torch.cuda.device_count()
mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
if __name__ == '__main__':
main()
在运行分布式训练脚本时，需要设置一些环境变量来指定进程数和GPU设备。例如：
export MASTER_ADDR='localhost'
export MASTER_PORT='12345'
python -m torch.distributed.launch --nproc_per_node=4 your_script.py
其中，--nproc_per_node参数指定了每个节点上的GPU数量。
通过以上步骤，你可以在CentOS上配置和使用PyTorch进行并行计算。根据你的具体需求选择合适的方法，并根据实际情况调整代码。
