Pytorch 多块GPU的使用详解

在PyTorch中，可以使用多块GPU来加速模型训练。以下是使用多块GPU的详细攻略：

1. 检查GPU是否可用

首先，需要检查GPU是否可用。可以使用以下代码检查GPU是否可用：

import torch

if torch.cuda.is_available():
    print('GPU is available!')
else:
    print('GPU is not available!')

如果输出结果为“GPU is available!”，则表示GPU可用。

1. 定义模型

接下来，需要定义模型。可以使用以下代码定义一个简单的模型：

import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, 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

在上面的代码中，定义了一个简单的模型，包含两个线性层和一个ReLU激活函数。

1. 将模型放到GPU上

接下来，需要将模型放到GPU上。可以使用以下代码将模型放到GPU上：

model = SimpleModel()
if torch.cuda.is_available():
    model.cuda()

在上面的代码中，如果GPU可用，则使用model.cuda()函数将模型放到GPU上。

1. 定义数据集和数据加载器

接下来，需要定义数据集和数据加载器。可以使用以下代码定义一个简单的数据集和数据加载器：

import torch.utils.data as data

class SimpleDataset(data.Dataset):
    def __init__(self):
        self.data = torch.randn(100, 10)
        self.labels = torch.randint(0, 2, (100,))

    def __getitem__(self, index):
        x = self.data[index]
        y = self.labels[index]
        return x, y

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

dataset = SimpleDataset()
dataloader = data.DataLoader(dataset, batch_size=10, shuffle=True)

在上面的代码中，定义了一个简单的数据集和数据加载器，包含100个样本和10个特征。

1. 定义优化器和损失函数

接下来，需要定义优化器和损失函数。可以使用以下代码定义一个简单的优化器和损失函数：

import torch.optim as optim

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

在上面的代码中，定义了一个简单的随机梯度下降优化器和交叉熵损失函数。

1. 训练模型

最后，可以使用以下代码训练模型：

for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(dataloader, 0):
        inputs, labels = data
        if torch.cuda.is_available():
            inputs = inputs.cuda()
            labels = labels.cuda()

        optimizer.zero_grad()

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

        running_loss += loss.item()
        if i % 10 == 9:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 10))
            running_loss = 0.0

在上面的代码中，使用一个简单的循环来训练模型。在每个epoch中，使用一个简单的循环来遍历数据加载器中的所有批次。在每个批次中，将输入和标签放到GPU上（如果GPU可用），然后使用优化器和损失函数来计算损失并更新模型参数。最后，输出每个epoch的平均损失。

以下是两个示例说明，用于说明如何在PyTorch中使用多块GPU：

示例1：使用DataParallel

可以使用DataParallel来自动将模型复制到多个GPU上，并将批次分配给不同的GPU。以下是使用DataParallel的示例代码：

import torch.nn as nn
import torch.nn.parallel

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, 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

model = SimpleModel()
if torch.cuda.is_available():
    model = nn.parallel.DataParallel(model)

# 训练模型

在上面的代码中，使用nn.parallel.DataParallel()函数将模型复制到多个GPU上。

示例2：手动分配批次

可以手动将批次分配给不同的GPU。以下是手动分配批次的示例代码：

import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, 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

model = SimpleModel()
if torch.cuda.is_available():
    model.cuda()

# 训练模型
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(dataloader, 0):
        inputs, labels = data
        if torch.cuda.is_available():
            inputs = inputs.cuda()
            labels = labels.cuda()

        # 手动分配批次
        inputs = torch.split(inputs, len(inputs) // torch.cuda.device_count())
        labels = torch.split(labels, len(labels) // torch.cuda.device_count())

        optimizer.zero_grad()

        for j in range(torch.cuda.device_count()):
            outputs = model(inputs[j])
            loss = criterion(outputs, labels[j])
            loss.backward()

        optimizer.step()

        running_loss += loss.item()
        if i % 10 == 9:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 10))
            running_loss = 0.0

在上面的代码中，使用torch.cuda.device_count()函数获取GPU数量，并手动将批次分配给不同的GPU。在每个批次中，将输入和标签分割成多个子批次，并在每个子批次上计算损失和梯度。最后，使用optimizer.step()函数来更新模型参数。

这是使用多块GPU的完整攻略，包括检查GPU是否可用、定义模型、将模型放到GPU上、定义数据集和数据加载器、定义优化器和损失函数以及训练模型的示例说明。同时，还包括使用DataParallel和手动分配批次的示例说明。希望对您有所帮助！