PyTorch 快速入门 (3)：图示分类器训练与多 GPU 训练

网络定义

对于这个图像分类任务，我们将使用一个简单的神经网络，称为 LeNet。这个网络由两个卷积层和两个全连接层组成。

import torch
import torch.nn as nn
import torch.nn.functional as F

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

数据处理和加载

接下来，我们需要准备 CIFAR10 数据集。这个数据集包含 60,000 张 32x32 的彩色图像，分为 10 个类别。我们可以使用 PyTorch 内置的 torchvision 库来下载和加载数据集。

import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                           download=True, transform=transform)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64,
                                         shuffle=True)

训练网络模型

现在，我们可以开始训练网络模型了。我们将使用交叉熵损失函数和随机梯度下降优化器来训练模型。

import torch.optim as optim

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

for epoch in range(20):
    for i, data in enumerate(train_loader):
        inputs, labels = data
        optimizer.zero_grad()

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

        optimizer.step()

    print(f'Epoch: {epoch + 1}, Loss: {loss.item()}')

测试模型性能

训练完成后，我们可以使用测试数据集来评估模型的性能。

test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                          download=True, transform=transform)

test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64,
                                        shuffle=True)

correct = 0
total = 0

with torch.no_grad():
    for data in test_loader:
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')

多 GPU 训练

如果有多个 GPU 可用，我们可以使用它们来并行训练模型，以提高训练速度。

import torch.distributed as dist
import torch.multiprocessing as mp

def train_with_multi_gpu(gpu, args):
    dist.init_process_group(backend='nccl', init_method='env://',
                            world_size=args.world_size, rank=gpu)

    model = LeNet()
    model.cuda(gpu)

    # 将模型分布到所有 GPU 上
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])

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

    for epoch in range(args.epochs):
        for i, data in enumerate(train_loader):
            inputs, labels = data
            inputs, labels = inputs.cuda(gpu), labels.cuda(gpu)
            optimizer.zero_grad()

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

            optimizer.step()

        if gpu == 0:
            print(f'Epoch: {epoch + 1}, Loss: {loss.item()}')

def main():
    args = parse_args()
    mp.spawn(train_with_multi_gpu,
             args=(args,),
             nprocs=args.world_size,
             join=True)

if __name__ == '__main__':
    main()

在运行脚本时，需要使用以下命令：

python main.py --world_size 2 --lr 0.01 --momentum 0.9 --epochs 20

其中，--world_size 指定了使用的 GPU 数量，--lr 和 --momentum 指定了优化器的学习率和动量，--epochs 指定了训练的轮数。

总结

在本教程中，我们介绍了如何使用 PyTorch 在 CIFAR10 数据集上训练一个图像分类器。我们讲解了网络定义、数据处理和加载、训练网络模型和测试模型性能的整个流程。我们还介绍了如何使用多 GPU 来训练网络模型，以提高训练速度。