pytorh使用多个GPU进行运算的代码

可以使用以下代码实现 PyTorch 在多个 GPU 上进行运算：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

# 设置使用的 GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.fc2(x)
        return x

# 定义数据集
class MyDataset(torch.utils.data.Dataset):
    def __init__(self):
        self.data = torch.randn(100, 10)
        self.label = torch.randint(0, 2, (100,))

    def __getitem__(self, index):
        return self.data[index], self.label[index]

    def __len__(self):
        return len(self.data)

# 定义训练函数
def train(model, dataloader, optimizer, criterion):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(dataloader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    return running_loss / len(dataloader)

# 定义测试函数
def test(model, dataloader, criterion):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for data in dataloader:
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)

            outputs = model(inputs)
            loss = criterion(outputs, labels)

            running_loss += loss.item()

            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    return running_loss / len(dataloader), correct / total

# 定义主函数
def main():
    # 初始化模型、数据集、优化器、损失函数
    model = MyModel()
    model = nn.DataParallel(model)
    model.to(device)

    dataset = MyDataset()
    dataloader = DataLoader(dataset, batch_size=10, shuffle=True)

    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    criterion = nn.CrossEntropyLoss()

    # 训练模型
    for epoch in range(10):
        train_loss = train(model, dataloader, optimizer, criterion)
        test_loss, test_acc = test(model, dataloader, criterion)
        print("Epoch %d, Train Loss: %.3f, Test Loss: %.3f, Test Acc: %.3f" % (epoch+1, train_loss, test_loss, test_acc))

if __name__ == '__main__':
    main()

以上代码中，我们使用了 `nn.DataParallel` 将模型并行在多个 GPU 上运算。在训练和测试函数中，我们将数据和模型都移动到 GPU 上进行运算。