文章目录
- 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.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在逐渐减小：
