优化器根据梯度调整参数,以达到降低误差
import torch.optim
import torchvision
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader
# 加载CIFAR10测试数据集,设置transform将图像转换为Tensor
dataset = torchvision.datasets.CIFAR10("./data", train=False, transform=torchvision.transforms.ToTensor(),
download=True)
# 创建数据加载器,设置批量大小为64
dataloader = DataLoader(dataset, batch_size=64)
# 定义卷积神经网络模型
class TY(nn.Module):
def __init__(self):
super(TY, self).__init__()
# 构建网络结构:3个卷积层+池化层组合,2个全连接层
self.model1 = Sequential(
Conv2d(3, 32, 5, padding=2), # 输入3通道,输出32通道,卷积核5x5
MaxPool2d(2), # 最大池化,步长2
Conv2d(32, 32, 5, padding=2), # 第二层卷积
MaxPool2d(2), # 第二次池化
Conv2d(32, 64, 5, padding=2), # 第三层卷积
MaxPool2d(2), # 第三次池化
Flatten(), # 将多维张量展平为向量
Linear(1024, 64), # 全连接层,输入1024维,输出64维
Linear(64, 10), # 输出层,10个类别对应10个输出
)
def forward(self, x):
# 定义前向传播路径
x = self.model1(x)
return x
# 定义损失函数(交叉熵损失适用于多分类问题)
loss = nn.CrossEntropyLoss()
# 实例化模型
ty = TY()
# 定义优化器(随机梯度下降),设置学习率为0.01
optim = torch.optim.SGD(ty.parameters(), lr=0.01)
# 训练20个完整轮次
for epoch in range(20):
running_loss = 0.0 # 初始化本轮累计损失
# 遍历数据加载器中的每个批次
for data in dataloader:
imgs, targets = data # 获取图像和标签
outputs = ty(imgs) # 前向传播
result_loss = loss(outputs, targets) # 计算损失
optim.zero_grad() # 梯度清零,防止累积
result_loss.backward() # 反向传播计算梯度
optim.step() # 更新模型参数
running_loss += result_loss # 累加损失值
# 打印本轮训练的累计损失
print(f"Epoch {epoch+1}, Loss: {running_loss}")