文章目录
- 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("------------------------------")
运行结果如下，发现在每一轮过后，Loss在逐渐减小：
)
)
