如何在Torch中实现多GPU训练

在Torch中实现多GPU训练可以通过使用nn.DataParallel模块来实现。nn.DataParallel模块可以将模型复制到多个GPU上，并自动将数据分发到每个GPU上进行计算，并最终将结果进行合并。以下是一个简单示例：

import torchimport torch.nn as nn# 定义模型class Model(nn.Module):    def __init__(self):        super(Model, self).__init__()        self.fc1 = nn.Linear(10, 5)        self.fc2 = nn.Linear(5, 1)        def forward(self, x):        x = self.fc1(x)        x = self.fc2(x)        return xmodel = Model()# 检查是否有多个GPU可用if torch.cuda.device_count() > 1:    model = nn.DataParallel(model)# 将模型加载到GPUdevice = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model = model.to(device)# 定义优化器和损失函数optimizer = torch.optim.SGD(model.parameters(), lr=0.001)criterion = nn.MSELoss()# 训练模型for epoch in range(num_epochs):    for inputs, targets in dataloader:        inputs, targets = inputs.to(device), targets.to(device)                optimizer.zero_grad()        outputs = model(inputs)        loss = criterion(outputs, targets)        loss.backward()        optimizer.step()

在上面的示例中，首先检查是否有多个GPU可用，如果有，则将模型放入nn.DataParallel中。然后将模型加载到GPU上进行训练。在训练过程中，数据会自动分发到每个GPU上，并在多个GPU上并行计算，从而加快训练速度。