在Linux环境下，PyTorch提供了多种并行计算的方法，主要包括数据并行和模型并行。以下是实现这些并行计算的步骤：
数据并行是指将数据分割成多个小批次，然后将这些小批次分配到不同的GPU上进行并行处理。
安装PyTorch： 确保你已经安装了支持GPU的PyTorch版本。可以使用以下命令安装：
pip install torch torchvision torchaudio 使用torch.nn.DataParallel： DataParallel是PyTorch中用于数据并行的模块。以下是一个简单的示例：
import torch import torch.nn as nn from torch.utils.data import DataLoader # 定义一个简单的模型 class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.fc = nn.Linear(10, 1) def forward(self, x): return self.fc(x) # 创建模型实例 model = SimpleModel() # 检查是否有可用的GPU if torch.cuda.device_count() > 1: print(f"Let's use {torch.cuda.device_count()} GPUs!") # 包装模型以使用多个GPU model = nn.DataParallel(model) # 将模型移动到GPU model.to('cuda') # 创建数据加载器 dataset = torch.randn(100, 10) # 示例数据 labels = torch.randn(100, 1) # 示例标签 dataloader = DataLoader(dataset, batch_size=10, shuffle=True) # 训练模型 for data, target in dataloader: data, target = data.to('cuda'), target.to('cuda') output = model(data) loss = nn.MSELoss()(output, target) loss.backward() optimizer.step() optimizer.zero_grad() 模型并行是指将模型的不同部分分配到不同的GPU上进行并行处理。
定义模型并分割： 假设我们有一个较大的模型，可以将其分割成多个部分，每个部分放在不同的GPU上。
import torch import torch.nn as nn class LargeModel(nn.Module): def __init__(self): super(LargeModel, self).__init__() self.part1 = nn.Linear(10, 10).to('cuda:0') self.part2 = nn.Linear(10, 1).to('cuda:1') def forward(self, x): x = x.to('cuda:0') x = self.part1(x) x = x.to('cuda:1') x = self.part2(x) return x model = LargeModel() 训练模型： 在训练过程中，需要手动管理数据在不同GPU之间的传递。
dataset = torch.randn(100, 10) # 示例数据 labels = torch.randn(100, 1) # 示例标签 dataloader = DataLoader(dataset, batch_size=10, shuffle=True) optimizer = torch.optim.SGD(model.parameters(), lr=0.01) for data, target in dataloader: data, target = data.to('cuda:0'), target.to('cuda:1') optimizer.zero_grad() output = model(data) loss = nn.MSELoss()(output, target) loss.backward() optimizer.step() DataParallel时，确保所有GPU上的模型参数是同步的。通过以上方法，你可以在Linux环境下实现PyTorch的并行计算，从而提高训练效率。
