Pytorch学习笔记——在GPU上进行训练

1. 环境准备

首先，确保PyTorch已经安装，且CUDA（NVIDIA的并行计算平台和编程模型）支持已经正确配置。可以通过以下代码检查CUDA是否可用：

print(torch.cuda.is_available())  # 如果返回True，则CUDA可用

配置PyTorch环境和CUDA支持，可以参考我写的这篇博客
Pytorch学习笔记——环境配置安装

2. 导入必要库

import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Linear, Flatten, Sequential
from torch.utils.data import DataLoader

• torch 是PyTorch的核心库。
• torchvision 提供了用于计算机视觉任务的工具，包括数据集和变换。
• nn 包含了构建神经网络所需的各种模块。
• DataLoader 用于加载数据集并进行批处理。

3. 加载数据集

使用CIFAR-10数据集进行训练，它是一个常用的小型图像数据集。加载数据集并创建数据加载器。

# 加载数据集
dataset = torchvision.datasets.CIFAR10(root="data1", train=False, transform=torchvision.transforms.ToTensor(), download=True)
# 创建数据加载器
dataloader = DataLoader(dataset, batch_size=64)

• torchvision.datasets.CIFAR10：下载并加载CIFAR-10数据集。
• DataLoader：将数据集划分为小批次，并进行数据加载。

4. 定义简单的神经网络模型

定义一个简单的卷积神经网络（CNN）模型：

self.model1 = Sequential(
            Conv2d(3, 32, 5, padding=2),  # 第一次卷积
            MaxPool2d(2),  # 第一次最大池化
            Conv2d(32, 32, 5, padding=2),  # 第二次卷积
            MaxPool2d(2),  # 第二次最大池化
            Conv2d(32, 64, 5, padding=2),  # 第三次卷积
            MaxPool2d(2),  # 第三次最大池化
            Flatten(),    # 展平层
            Linear(1024, 64),  # 第一个全连接层
            Linear(64, 10),  # 第二个全连接层
        )
def forward(self, x):
        x = self.model1(x)
        return x

• Conv2d：二维卷积层，用于提取图像特征。
• MaxPool2d：最大池化层，用于下采样。
• Flatten：将多维输入展平为一维，用于全连接层的输入。
• Linear：全连接层，用于分类任务。

5. 检查和设置GPU设备

需要检查是否有可用的GPU，并将模型和数据移动到GPU上。

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将模型和数据转移到GPU上
mynn = NN().to(device)
print(mynn)

• torch.device("cuda")：如果CUDA可用，则使用GPU；否则使用CPU。
• to(device)：将模型转移到指定的设备（CPU或GPU）。

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

定义损失函数和优化器来训练模型：

# 定义损失函数
loss = nn.CrossEntropyLoss().to(device)

# 定义优化器
optim = torch.optim.SGD(mynn.parameters(), lr=0.01)

• nn.CrossEntropyLoss：适用于分类问题的损失函数。
• torch.optim.SGD：随机梯度下降优化器。

7. 训练模型

通过多个epoch对模型进行训练。在每个epoch中，进行前向传播、计算损失、反向传播和参数更新：

# 多轮学习 0 - 20 20轮
for epoch in range(20):
    running_loss = 0.0
    for data in dataloader:
        # 确保数据也转移到GPU上
        imgs, targets = data[0].to(device), data[1].to(device)

        optim.zero_grad()  # 清零梯度缓存
        outputs = mynn(imgs)  # 前向传播
        loss_value = loss(outputs, targets)  # 计算损失
        loss_value.backward()  # 反向传播，计算梯度
        optim.step()  # 根据梯度更新权重

        running_loss += loss_value.item()  # 累加损失值

    print(f"Epoch {epoch + 1}, Loss: {running_loss / len(dataloader)}")
    print("------------------------------")

• optim.zero_grad()：清零之前计算的梯度。
• outputs = mynn(imgs)：进行前向传播。
• loss_value.backward()：进行反向传播，计算梯度。
• optim.step()：根据计算的梯度更新权重。
• running_loss：累加损失值以计算平均损失。

8. 全部代码展示及运行结果

# -*- coding: utf-8 -*-
# @Author: kk
import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Linear, Flatten, Sequential
from torch.utils.data import DataLoader

# 加载数据集
dataset = torchvision.datasets.CIFAR10(root="data1", train=False, transform=torchvision.transforms.ToTensor(), download=True)
# loader加载
dataloader = DataLoader(dataset, batch_size=64)

# 网络
class NN(nn.Module):
    def __init__(self):
        super(NN, self).__init__()
        self.model1 = Sequential(
            Conv2d(3, 32, 5, padding=2),  # 第一次卷积
            MaxPool2d(2),  # 第一次最大池化
            Conv2d(32, 32, 5, padding=2),  # 第二次卷积
            MaxPool2d(2),  # 第二次最大池化
            Conv2d(32, 64, 5, padding=2),  # 第三次卷积
            MaxPool2d(2),  # 第三次最大池化
            Flatten(),    # 展平层
            Linear(1024, 64),  # 第一个全连接层
            Linear(64, 10),  # 第二个全连接层
        )

    def forward(self, x):
        x = self.model1(x)
        return x

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将模型和数据转移到GPU上
mynn = NN().to(device)
print(mynn)
loss = nn.CrossEntropyLoss().to(device)

# 优化器
optim = torch.optim.SGD(mynn.parameters(), lr=0.01)

# 多轮学习  0 - 20  20轮
for epoch in range(20):
    running_loss = 0.0
    for data in dataloader:
        # 确保数据也转移到GPU上
        imgs, targets = data[0].to(device), data[1].to(device)

        optim.zero_grad()  # 清零梯度缓存
        outputs = mynn(imgs)  # 前向传播
        loss_value = loss(outputs, targets)  # 计算损失
        loss_value.backward()  # 反向传播，计算梯度
        optim.step()  # 根据梯度更新权重

        running_loss += loss_value.item()  # 累加损失值

    print(f"Epoch {epoch + 1}, Loss: {running_loss / len(dataloader)}")
    print("------------------------------")

运行结果如下，发现在每一轮过后，Loss在逐渐减小：