PyTorch支持使用多张显卡进行并行训练，可以使用`torch.nn.DataParallel`或`torch.nn.parallel.DistributedDataParallel`来实现。以下是使用`torch.nn.DataParallel`的示例代码： ```python import torch import torch.nn as nn from torch.utils.data import DataLoader # 定义模型 class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.fc1 = nn.Linear(10, 5) self.fc2 = nn.Linear(5, 1) def forward(self, x): x = self.fc1(x) x = torch.relu(x) x = self.fc2(x) return x # 定义数据集 class MyDataset(torch.utils.data.Dataset): def __init__(self): self.data = torch.randn(100, 10) self.targets = torch.randn(100, 1) def __getitem__(self, index): return self.data[index], self.targets[index] def __len__(self): return len(self.data) # 定义训练函数 def train(model, dataloader, optimizer, criterion): model.train() for i, (inputs, targets) in enumerate(dataloader): optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.step() # 创建模型、数据集、数据加载器、优化器、损失函数 model = MyModel() dataset = MyDataset() dataloader = DataLoader(dataset, batch_size=8, shuffle=True) optimizer = torch.optim.SGD(model.parameters(), lr=0.01) criterion = nn.MSELoss() # 使用DataParallel进行多GPU并行训练 model = nn.DataParallel(model) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.to(device) for epoch in range(10): train(model, dataloader, optimizer, criterion) ``` 在上面的代码中，我们首先定义了一个模型`MyModel`和一个数据集`MyDataset`。然后，我们使用`DataLoader`将数据集加载到内存中。接下来，我们创建了一个优化器和一个损失函数。最后，我们将模型移到GPU上，并使用`DataParallel`对其进行并行处理。在训练循环中，我们调用`train`函数来训练模型。`train`函数的参数分别是模型、数据加载器、优化器和损失函数。由于我们在模型上调用了`DataParallel`，因此在训练循环中，我们不需要手动处理多个GPU的并行计算。
