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

文章目录

- - 1. 环境准备
  - 2. 导入必要库
  - 3. 加载数据集
  - 4. 定义简单的神经网络模型
  - 5. 检查和设置GPU设备
  - 6. 定义损失函数和优化器
  - 7. 训练模型
  - 8. 全部代码展示及运行结果

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.0for 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.0for 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("------------------------------")