B站小土堆-pytorch深度学习快速入门笔记

看月亮的方源2025-09-04 21:01

部分函数说明

1. 卷积层

  1. 卷积层的作用
  • 提取输入图片中的信息，这些信息被称为图像特征，这些特征是由图像中的每个像素通过组合或者独立的方式所体现，比如图片的纹理特征，颜色特征。
  1. 这里的卷积操作是通过卷积核对每个通道的矩阵从左到右（卷积核一般是3x3的矩阵）从上至下进行互相关运算（先是从左到右，再是从上至下，所以卷积操作也会保留位置信息），就像一个小的窗口一样，从左上角一步步滑动到右下角，滑动的步长是个超参数，互相关运算的意思就是对应位置相乘再相加，最后把三个通道的值也对应加起来得到一个值。
py 复制代码 
# 导入必要的库
import torch  # PyTorch深度学习框架主库
import torchvision  # 主要处理图像数据，包含一些常用的数据集、模型、转换函数等
from torch import nn  # 神经网络模块
from torch.nn import Conv2d  # 二维卷积层
from torch.utils.data import DataLoader  # 数据加载器，用于批量加载数据
from torch.utils.tensorboard import SummaryWriter  # TensorBoard可视化工具

# 加载CIFAR10数据集
dataset = torchvision.datasets.CIFAR10(
    "../data",  # 数据集存储路径
    train=False,  # 使用测试集（而非训练集）
    transform=torchvision.transforms.ToTensor(),  # 将图像转换为Tensor格式
    download=True  # 如果数据集不存在则自动下载
)

# 创建数据加载器，用于批量获取数据
dataloader = DataLoader(
    dataset,  # 使用的数据集
    batch_size=64  # 每批加载64个样本
)


# 定义神经网络类
class TuiDui(nn.Module):
    def __init__(self):
        super(TuiDui, self).__init__()  # 调用父类nn.Module的初始化方法
        # 定义卷积层
        self.conv1 = Conv2d(
            in_channels=3,  # 输入通道数（RGB图像为3通道）
            out_channels=6,  # 输出通道数（卷积核数量）
            kernel_size=3,  # 卷积核大小（3x3）
            stride=1,  # 卷积步长
            padding=0  # 填充大小（无填充）
        )

    def forward(self, x):
        x = self.conv1(x)  # 前向传播：对输入x进行卷积操作
        return x


# 创建神经网络实例
tudui = TuiDui()

# 创建TensorBoard SummaryWriter对象，用于记录日志
writer = SummaryWriter("./juanlogs")

# 初始化步数计数器
step = 0

# 遍历数据加载器中的每个批次
for data in dataloader:
    imgs, targets = data  # 解包数据批次为图像和标签
    output = tudui(imgs)  # 将图像输入网络获取输出

    # 记录输入图像到TensorBoard
    writer.add_images("input", imgs, step)

    # 重塑输出张量形状以便可视化
    # 将[64, 6, 30, 30]重塑为[128, 3, 30, 30]（因为6个通道可以分成2组3通道图像）
    output = torch.reshape(output, (-1, 3, 30, 30))

    # 记录输出图像到TensorBoard
    writer.add_images("output", output, step)

    # 更新步数计数器
    step += 1

2. 池化层的作用

  1. 池化层的作用是对卷积层中提取的特征进行挑选。主要有以下几个作用：

  • 挑选不受位置干扰的图像信息。

  • 对特征进行降维，提高后续特征的感受野，也就是让池化后的一个像素对应前面图片中的一个区域。

  • 因为池化层是不进行反向传播的，而且池化层减少了特征图的变量个数，所以池化层可以减少计算量。

  1. 常见的池化操作有最大池化和平均池化，池化层是由n×n大小的矩阵窗口滑动来进行计算的，类似于卷积层，只不过不是做互相关运算，而是求n×n大小的矩阵中的最大值、平均值等。

2.1 最大池化的使用

python 复制代码 
# 导入PyTorch深度学习框架
import torch
# 导入torchvision中的数据集模块
import torchvision.datasets
# 从PyTorch中导入神经网络模块
from torch import nn
# 导入数据加载器，用于批量加载数据
from torch.utils.data import DataLoader
# 导入TensorBoard可视化工具
from torch.utils.tensorboard import SummaryWriter

# 加载CIFAR10数据集
dataset = torchvision.datasets.CIFAR10(
    root='./data',  # 数据集存储的根目录
    train=False,    # 使用测试集（而不是训练集）
    download=True,  # 如果数据集不存在则自动下载
    # 数据转换：将PIL图像转换为Tensor格式，并自动归一化到[0,1]范围
    transform=torchvision.transforms.ToTensor()
)

# 创建数据加载器，用于批量获取数据
dataloader = DataLoader(dataset, batch_size=64)  # 每批加载64个样本


# 定义神经网络类
class TuDui(nn.Module):
    def __init__(self):
        # 调用父类构造函数进行初始化
        super(TuDui, self).__init__()
        # 定义最大池化层
        self.maxpool = nn.MaxPool2d(
            kernel_size=3,   # 池化窗口大小（3x3）
            ceil_mode=True   # 当池化窗口无法完全覆盖输入时，保留剩余部分（向上取整）
        )
        # ceil_mode=True，最大池化，保留输入的特征（即使不能完全被窗口覆盖）

    def forward(self, input):
        # 前向传播：对输入进行最大池化操作
        output = self.maxpool(input)
        return output


# 创建神经网络实例
tudui = TuDui()
# 初始化步数计数器
step = 0
# 创建TensorBoard日志写入器，指定日志保存路径
writer = SummaryWriter("./logs_max_pool")

# 遍历数据加载器中的所有批次
for data in dataloader:
    # 解包数据，获取图像和对应的标签
    imgs, targets = data
    # 将输入图像写入TensorBoard
    writer.add_images("input", imgs, step)
    # 将图像输入网络，获取输出（经过最大池化后的特征图）
    output = tudui(imgs)
    # 将输出图像写入TensorBoard
    writer.add_images("output", output, step)
    # 增加步数计数器
    step += 1

# 关闭TensorBoard写入器，确保所有数据已写入磁盘
writer.close()

保留输入的特征 代码运行后的效果如下：

3. 非线性激活

非线性激活，相当于两层神经元之间的关系函数，上层的输出，被激活函数作用得到下层的输入。

python 复制代码 
# 导入PyTorch深度学习框架
import torch
# 导入TorchVision库，提供计算机视觉相关的数据集和转换
import torchvision.datasets
# 从PyTorch中导入神经网络模块
from torch import nn
# 从神经网络模块中导入ReLU激活函数
from torch.nn import ReLU
# 导入数据加载器，用于批量加载数据
from torch.utils.data import DataLoader
# 导入TensorBoard日志记录工具
from torch.utils.tensorboard import SummaryWriter
# 导入urllib3的文件上传模块（此导入可能多余，因为下面使用的是SummaryWriter）
from urllib3.filepost import writer

# 创建一个2x2的张量，包含4个数值
input = torch.tensor([[1, -0.5],
                      [-1, 3]])
# 将张量重塑为4维形状，适合作为卷积神经网络的输入
# 形状变为(batch_size=1, channels=1, height=2, width=2)
input = torch.reshape(input, (-1,  1, 2, 2))
# 打印重塑后的张量形状
print(input.shape)

# 加载CIFAR-10数据集
dataset = torchvision.datasets.CIFAR10(root='./data',  # 数据集存储路径
                                       train=True,     # 加载训练集
                                       download=True,  # 如果数据集不存在则下载
                                       # 将PIL图像转换为Tensor，并归一化到[0,1]范围
                                       transform=torchvision.transforms.ToTensor())
# 创建数据加载器，用于批量加载数据
dataloader = DataLoader(dataset, batch_size=64)  # 每批加载64个样本

# 定义一个名为TuDui的神经网络类，继承自nn.Module
class TuDui(nn.Module):
    # 类的初始化函数
    def __init__(self):
        # 调用父类的初始化方法
        super().__init__()
        # 创建一个ReLU激活函数实例（虽然定义了但未在前向传播中使用）
        self.relu = ReLU(inplace=False)
        # 创建一个Sigmoid激活函数实例
        self.sigmoid1 = nn.Sigmoid()

    # 定义前向传播过程
    def forward(self, input):
        # 对输入应用Sigmoid激活函数（将值压缩到0-1范围内）
        output = self.sigmoid1(input)
        # 返回处理后的输出
        return output

# 创建TuDui类的实例
tudui = TuDui()
# 初始化步数计数器，用于TensorBoard记录
step = 0
# 创建SummaryWriter实例，指定日志保存路径
writer = SummaryWriter("./relu_logs")

# 遍历数据加载器中的所有批次
for data in dataloader:
    # 解包数据，获取图像和对应的标签
    imgs, targets = data
    # 将输入图像添加到TensorBoard中，用于可视化
    writer.add_images("input", imgs, step)
    # 将图像通过神经网络模型进行处理
    outputs = tudui(imgs)
    # 将处理后的输出图像添加到TensorBoard中，用于可视化
    writer.add_images("output", outputs, step)
    # 增加步数计数器
    step += 1

# 关闭SummaryWriter，确保所有数据都已写入磁盘
writer.close()

4. 小实战_sequential

python 复制代码 
import torch
from torch import nn
from torch.utils.tensorboard import SummaryWriter
from urllib3.filepost import writer


class TuiDui(nn.Module):
    def __init__(self):
        super().__init__()
        # self.conv1 = nn.Conv2d(3, 32, 5, padding=2)
        # self.maxpool = nn.MaxPool2d(2)
        # self.conv2 = nn.Conv2d(32, 32, 5, padding=2)
        # self.maxpool2 = nn.MaxPool2d(2)
        # self.conv3 = nn.Conv2d(32, 64, 5, padding=2)
        # self.maxpool3 = nn.MaxPool2d(2)
        # self.flatten = nn.Flatten()
        # self.linear1 = nn.Linear(1024, 64)
        # self.linear2 = nn.Linear(64, 10)
        # Sequential 等价于上面的写法
        # 使用nn.Sequential顺序容器构建网络模型，替代之前注释掉的逐层定义方式
        self.model1 = nn.Sequential(
            # 第一个卷积层：输入通道3（RGB图像），输出通道32，卷积核大小5，填充2（保持特征图尺寸）
            nn.Conv2d(3, 32, 5, padding=2),
            # 最大池化层，池化窗口大小为2（下采样因子为2）
            nn.MaxPool2d(2),
            # 第二个卷积层：输入通道32，输出通道32，卷积核大小5，填充2
            nn.Conv2d(32, 32, 5, padding=2),
            # 第二个最大池化层
            nn.MaxPool2d(2),
            # 第三个卷积层：输入通道32，输出通道64，卷积核大小5，填充2
            nn.Conv2d(32, 64, 5, padding=2),
            # 第三个最大池化层
            nn.MaxPool2d(2),
            # 展平层：将多维特征张量转换为一维，便于全连接层处理
            nn.Flatten(),
            # 第一个全连接层：输入特征数1024（64*4*4=1024），输出64
            nn.Linear(1024, 64),
            # 第二个全连接层（输出层）：输入64，输出10（对应10个类别）
            nn.Linear(64, 10)
        )

  
    def forward(self, x):
        # x = self.conv1(x)
        # x = self.maxpool(x)
        # x = self.conv2(x)
        # x = self.maxpool2(x)
        # x = self.conv3(x)
        # x = self.maxpool3(x)
        # x = self.flatten(x)
        # x = self.linear1(x)
        # x = self.linear2(x)
        x = self.model1(x)
        return x


# 创建TuiDui模型实例
tudui = TuiDui()
# 打印模型结构
print(tudui)

# 创建模拟输入数据：批次大小64，3通道，32x32像素的图像（全部初始化为1）
input = torch.ones(64, 3, 32, 32)
# 将输入数据通过模型前向传播得到输出
output = tudui(input)
# 打印输出张量的形状
print(output.shape)

# 创建TensorBoard的SummaryWriter对象，日志保存到"./log_seq"目录
writer = SummaryWriter("./log_seq")
# 将模型结构和示例输入添加到TensorBoard日志
writer.add_graph(tudui, input)
# 关闭SummaryWriter
writer.close()

效果如图所示

5. 损失函数

python 复制代码 
# 导入PyTorch深度学习框架
import torch
# 从PyTorch的神经网络模块中导入L1Loss（平均绝对误差损失）
from torch.nn import L1Loss
# 导入PyTorch的神经网络模块
from torch import nn

# 创建输入张量，包含数值1, 2, 3，并指定数据类型为浮点型
inputs = torch.tensor([1, 2, 3], dtype=torch.float)
# 创建目标张量，包含数值1, 2, 5，并指定数据类型为浮点型
targets = torch.tensor([1, 2, 5], dtype=torch.float)

# 将输入张量从形状[3]重塑为[1, 1, 1, 3]
# 这增加了批次、通道等维度，使其符合深度学习模型的输入格式要求
inputs = torch.reshape(inputs, (1, 1, 1, 3))
# 同样重塑目标张量的形状，使其与输入张量形状一致
targets = torch.reshape(targets, (1, 1, 1, 3))

# 创建L1损失函数实例，设置reduction参数为'sum'表示计算所有元素的绝对误差之和
loss = L1Loss(reduction='sum')
# 计算输入和目标之间的L1损失
result = loss(inputs, targets)

# 创建均方误差(MSE)损失函数实例
loss_me = nn.MSELoss()
# 计算输入和目标之间的均方误差损失
result_mse = loss_me(inputs, targets)

# 打印L1损失结果
print(result)
# 打印均方误差损失结果
print(result_mse)

代码作用说明

这段代码演示了PyTorch中两种常见损失函数的计算：

  1. L1损失（平均绝对误差）：计算预测值与目标值之间绝对差异的总和

    • 输入值：1, 2, 3
    • 目标值：1, 2, 5
    • 绝对误差：|1-1| + |2-2| + |3-5| = 0 + 0 + 2 = 2

  2. MSE损失（均方误差）：计算预测值与目标值之间平方差异的平均值

    • 输入值：1, 2, 3
    • 目标值：1, 2, 5
    • 平方误差：(1-1)² + (2-2)² + (3-5)² = 0 + 0 + 4 = 4
    • 平均值：4 / 3 ≈ 1.333

6.前向传播和反向传播（梯度计算）

python 复制代码 
# 导入PyTorch计算机视觉库
import torchvision
# 导入PyTorch神经网络模块
from torch import nn
# 导入数据加载器，用于批量加载数据
from torch.utils.data import DataLoader

# 加载CIFAR-10数据集
dataset = torchvision.datasets.CIFAR10(
    "./data",  # 数据集存储路径
    train=False,  # 使用测试集（非训练集）
    transform=torchvision.transforms.ToTensor(),  # 将图像转换为Tensor格式
    download=True  # 如果数据集不存在则自动下载
)

# 创建数据加载器，设置批量大小为64
dataloader = DataLoader(dataset, batch_size=64)

# 定义神经网络模型类
class TuiDui(nn.Module):
    def __init__(self):
        # 调用父类构造函数
        super().__init__()
        
        # 使用nn.Sequential顺序容器构建网络模型
        self.model1 = nn.Sequential(
            # 第一个卷积层：输入通道3（RGB图像），输出通道32，卷积核大小5，填充2（保持特征图尺寸）
            nn.Conv2d(3, 32, 5, padding=2),
            # 最大池化层，池化窗口大小为2（下采样因子为2）
            nn.MaxPool2d(2),
            # 第二个卷积层：输入通道32，输出通道32，卷积核大小5，填充2
            nn.Conv2d(32, 32, 5, padding=2),
            # 第二个最大池化层
            nn.MaxPool2d(2),
            # 第三个卷积层：输入通道32，输出通道64，卷积核大小5，填充2
            nn.Conv2d(32, 64, 5, padding=2),
            # 第三个最大池化层
            nn.MaxPool2d(2),
            # 展平层：将多维特征张量转换为一维，便于全连接层处理
            nn.Flatten(),
            # 第一个全连接层：输入特征数1024（64*4*4=1024），输出64
            nn.Linear(1024, 64),
            # 第二个全连接层（输出层）：输入64，输出10（对应CIFAR-10的10个类别）
            nn.Linear(64, 10)
        )
    
    # 定义前向传播过程
    def forward(self, x):
        x = self.model1(x)
        return x

# 实例化神经网络模型
tudui = TuiDui()

# 定义交叉熵损失函数，适用于多分类问题
loss = nn.CrossEntropyLoss()

# 遍历数据加载器中的每个批次
for data in dataloader:
    # 解包数据，获取图像和对应的标签
    imgs, targets = data
    # 将图像输入网络，获取预测输出
    outputs = tudui(imgs)
    # 计算预测输出与真实标签之间的损失
    result_loss = loss(outputs, targets)
    # 反向传播，计算梯度
    result_loss.backward()
    # 打印提示信息，表示一次前向传播和反向传播完成
    print("ok")

代码功能说明

这段代码实现了一个完整的深度学习训练流程，主要功能包括：

  1. 数据准备：

    • 加载CIFAR-10数据集（10个类别的彩色图像数据集）
    • 将图像转换为PyTorch Tensor格式
    • 创建数据加载器，实现批量数据加载

  2. 神经网络模型定义：

    • 创建了一个名为"TuiDui"的卷积神经网络
    • 网络结构包含3个卷积层和2个全连接层
    • 使用最大池化层进行下采样
    • 最终输出10个节点，对应CIFAR-10的10个类别

  3. 训练流程：

    • 实例化模型和损失函数
    • 遍历数据集中的每个批次
    • 执行前向传播计算预测值
    • 计算损失值
    • 执行反向传播计算梯度

需要注意的是，这段代码只完成了前向传播和反向传播（梯度计算），但没有包含优化器步骤（权重更新），因此实际上并没有进行模型训练，只是演示了训练流程的前半部分。要完成完整的训练，还需要添加优化器和权重更新步骤。

7. 优化器

ps: 没搞明白，二刷的时候，再解决这个问题

python 复制代码 
import torch
import torchvision
# 导入PyTorch神经网络模块
from torch import nn
# 导入数据加载器，用于批量加载数据
from torch.utils.data import DataLoader

# 加载CIFAR-10数据集
dataset = torchvision.datasets.CIFAR10(
    "./data",  # 数据集存储路径
    train=False,  # 使用测试集（非训练集）
    transform=torchvision.transforms.ToTensor(),  # 将图像转换为Tensor格式
    download=True  # 如果数据集不存在则自动下载
)

# 创建数据加载器，设置批量大小为64
dataloader = DataLoader(dataset, batch_size=64)


# 定义神经网络模型类
class TuiDui(nn.Module):
    def __init__(self):
        # 调用父类构造函数
        super().__init__()

        # 使用nn.Sequential顺序容器构建网络模型
        self.model1 = nn.Sequential(
            # 第一个卷积层：输入通道3（RGB图像），输出通道32，卷积核大小5，填充2（保持特征图尺寸）
            nn.Conv2d(3, 32, 5, padding=2),
            # 最大池化层，池化窗口大小为2（下采样因子为2）
            nn.MaxPool2d(2),
            # 第二个卷积层：输入通道32，输出通道32，卷积核大小5，填充2
            nn.Conv2d(32, 32, 5, padding=2),
            # 第二个最大池化层
            nn.MaxPool2d(2),
            # 第三个卷积层：输入通道32，输出通道64，卷积核大小5，填充2
            nn.Conv2d(32, 64, 5, padding=2),
            # 第三个最大池化层
            nn.MaxPool2d(2),
            # 展平层：将多维特征张量转换为一维，便于全连接层处理
            nn.Flatten(),
            # 第一个全连接层：输入特征数1024（64*4*4=1024），输出64
            nn.Linear(1024, 64),
            # 第二个全连接层（输出层）：输入64，输出10（对应CIFAR-10的10个类别）
            nn.Linear(64, 10)
        )

    # 定义前向传播过程
    def forward(self, x):
        x = self.model1(x)
        return x


# 定义交叉熵损失函数，适用于多分类问题
loss = nn.CrossEntropyLoss()
# 实例化神经网络模型
tudui = TuiDui()
# 定义优化器
optim = torch.optim.SGD(tudui.parameters(), lr=0.01)

# 遍历数据加载器中的每个批次
for data in dataloader:
    # 解包数据，获取图像和对应的标签
    imgs, targets = data
    # 将图像输入网络，获取预测输出
    outputs = tudui(imgs)
    # 计算预测输出与真实标签之间的损失
    result_loss = loss(outputs, targets)
    # 优化器梯度清零
    optim.zero_grad()
    # 调用损失函数 反向传播，计算梯度
    result_loss.backward()
    # 对每个模型的参数进行调优
    optim.step()

8. 网络模型的保存与读取

网络模型的保存

python 复制代码 
# 导入PyTorch深度学习框架
import torch
# 导入torchvision计算机视觉库，包含常用模型、数据集和图像处理工具
import torchvision

# 创建一个VGG16模型实例，不加载预训练权重
# weights=None 表示使用随机初始化的权重，而不是预训练权重
# 这通常用于从头开始训练模型或加载自定义训练的权重
vgg16 = torchvision.models.vgg16(weights=None)

# 模型保存方式一：保存完整的模型（包括模型结构和所有参数）
# 这种方式会保存整个模型对象，包括网络架构和权重
# 优点是加载方便，直接使用torch.load()即可
# 缺点是文件较大，且可能与特定PyTorch版本绑定
torch.save(vgg16, 'vgg16_method1.pth')

# 模型保存方式二：只保存模型参数（官方推荐方式）
# vgg16.state_dict()返回一个包含模型所有参数的字典
# 这种方式只保存模型的权重，不保存网络结构
# 优点是文件较小，更加灵活，可以跨版本使用
# 缺点是加载时需要先创建模型结构，再加载参数
torch.save(vgg16.state_dict(), 'vgg16_method2.pth')

功能说明

保存PyTorch模型的方法：

  1. 方式一：保存完整模型

    • 使用torch.save(model, path)保存整个模型对象
    • 包含模型结构和所有参数
    • 加载时直接使用model = torch.load(path)

  2. 方式二：只保存模型参数（官方推荐）

    • 使用model.state_dict()获取模型参数字典
    • 使用torch.save(model.state_dict(), path)只保存参数
    • 加载时需要先创建模型结构，再使用model.load_state_dict(torch.load(path))

第二种方式是PyTorch官方推荐的，因为它更加灵活且文件更小，同时减少了与特定PyTorch版本的依赖关系。

网络模型的读取

python 复制代码 
import torch  # 导入PyTorch深度学习框架
import torchvision  # 导入torchvision计算机视觉库
from torchvision.models import VGG16_Weights  # 导入VGG16预训练权重

# 方式一：直接加载整个模型（包括模型结构和参数）
# 功能：从文件加载完整的模型，包含网络结构和训练好的权重
model = torch.load('vgg16_method1.pth', weights_only=False)
# 参数说明：
# 'vgg16_method1.pth' - 模型文件路径
# weights_only=False - 允许加载包含完整模型而不仅仅是权重的文件

# 方式二：分别加载模型结构和权重
# 功能：创建一个VGG16模型实例，不加载预训练权重（weights=None）
vgg16 = torchvision.models.vgg16(weights=None)
# 功能：加载仅包含模型权重的文件，并将其加载到模型结构中
vgg16.load_state_dict(torch.load('vgg16_method2.pth', weights_only=False))
# 参数说明：
# 'vgg16_method2.pth' - 仅包含模型状态字典（权重）的文件
# weights_only=False - 允许加载包含完整模型而不仅仅是权重的文件

# 功能：打印模型结构，显示所有层和参数
print(vgg16)

功能说明

这段代码展示了两种加载预训练VGG16模型的方法：

  1. 方式一：直接加载完整模型

    • 使用torch.load()直接加载包含模型结构和参数的文件
    • 适用于保存时使用了torch.save(model, path)的情况

  2. 方式二：分别加载模型结构和参数

    • 先创建模型结构实例，然后加载仅包含参数的状态字典
    • 适用于保存时使用了torch.save(model.state_dict(), path)的情况
    • 这种方式更灵活，可以在加载权重前修改模型结构

两种方法的主要区别在于模型文件的保存方式，方式一保存完整模型，方式二只保存模型参数。

9. 完整模型训练的套路

python 复制代码 
import torch
import torchvision  # PyTorch的计算机视觉库，提供数据集、模型架构和图像转换工具

# 导入数据加载工具和TensorBoard可视化工具
from torch.utils.data import DataLoader  # 用于批量加载数据
from torch.utils.tensorboard import SummaryWriter  # 用于训练过程可视化

# 导入自定义模型（假设存在model.py文件，其中包含模型定义）
from model import *

# 准备训练数据集（CIFAR-10）
train_data = torchvision.datasets.CIFAR10(
    root='./data',  # 数据集存储路径
    train=True,  # 加载训练集
    transform=torchvision.transforms.ToTensor(),  # 将PIL图像转换为PyTorch Tensor格式，并归一化到[0,1]
    download=True  # 如果本地不存在则自动下载
)

# 准备测试数据集
test_data = torchvision.datasets.CIFAR10(
    root='./data',
    train=False,  # 加载测试集
    transform=torchvision.transforms.ToTensor(),
    download=True
)

# 获取数据集尺寸
train_data_size = len(train_data)  # 训练集样本数量（50000）
test_data_size = len(test_data)  # 测试集样本数量（10000）

# 创建数据加载器（自动分批、打乱数据）
train_dataloader = DataLoader(train_data, batch_size=64)  # 每批64张图片，默认会打乱数据
test_dataloader = DataLoader(test_data, batch_size=64)  # 测试集不需要打乱

# 初始化自定义神经网络模型（假设model.py中存在TuiDui类）
tudui = TuiDui()

# 定义损失函数（交叉熵损失，适用于多分类问题）
loss_fn = nn.CrossEntropyLoss()

# 配置优化器（随机梯度下降）
learning_rate = 0.01
optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate)  # 使用SGD优化器，传入模型参数和学习率

# 初始化训练参数
total_train_step = 0  # 记录训练总步数（每个batch算一步）
total_test_step = 0  # 记录测试总次数（每个epoch测试一次）
epoch = 10  # 训练轮数

# 初始化TensorBoard日志记录器
writer = SummaryWriter(log_dir='./logs_train')  # 创建SummaryWriter对象，指定日志保存路径

# 开始训练循环
for i in range(epoch):
    print(f"----第{i + 1}轮训练开始---")

    # 训练阶段
    tudui.train()  # 设置模型为训练模式（启用dropout和batch normalization）
    for data in train_dataloader:
        imgs, targets = data  # 获取批量图像和对应标签
        outputs = tudui(imgs)  # 前向传播，计算模型输出
        loss = loss_fn(outputs, targets)  # 计算损失值

        # 反向传播优化
        optimizer.zero_grad()  # 清零梯度，防止梯度累积
        loss.backward()  # 反向传播，计算梯度
        optimizer.step()  # 更新模型参数

        # 记录训练信息
        total_train_step += 1
        if total_train_step % 100 == 0:  # 每100步输出一次日志
            print("训练次数：{}，loss:{}".format(total_train_step, loss.item()))
            writer.add_scalar('train_loss', loss.item(), total_train_step)  # 记录训练损失到TensorBoard

    # 测试阶段
    tudui.eval()  # 设置模型为评估模式（禁用dropout和batch normalization）
    total_test_loss = 0  # 累计测试损失
    total_accuracy = 0  # 累计正确预测的数量
    with torch.no_grad():  # 关闭梯度计算，节省内存和计算资源
        for data in test_dataloader:
            imgs, targets = data
            outputs = tudui(imgs)
            loss = loss_fn(outputs, targets)
            total_test_loss += loss.item()  # 累加测试损失
            
            # 计算准确率：找出输出中最大概率的类别，与真实标签比较
            accuracy = (outputs.argmax(1) == targets).sum()  # 计算当前batch中正确预测的数量
            total_accuracy += accuracy.item()  # 累加正确预测的数量
    
    # 记录测试结果
    print("整体测试集上的loss:{}".format(total_test_loss))
    print("整体测试集上的正确率：{}".format(total_accuracy / test_data_size))  # 计算整体准确率
    writer.add_scalar('test_loss', total_test_loss, total_test_step)  # 记录测试损失到TensorBoard
    writer.add_scalar('test_accuracy', total_accuracy / test_data_size, total_test_step)  # 记录测试准确率到TensorBoard
    total_test_step += 1

    # 保存模型 checkpoint
    torch.save(tudui, 'tudui_{}.pth'.format(i))  # 保存整个模型（包括结构和参数）
    # 模型保存方式二（注释掉的代码）
    # torch.save(tudui.state_dict(), 'tudui_{}.pth'.format(i))  # 只保存模型参数（更轻量）
    print("模型已保存")

# 关闭TensorBoard记录器
writer.close()

主要功能说明：

  1. 数据准备与加载： 使用CIFAR-10数据集（10个类别的60,000张32x32彩色图像）, 自动下载数据集并转换为PyTorch Tensor格式, 创建数据加载器实现批量处理（每批64张图像）。

  2. 模型训练： 使用自定义神经网络模型（TuiDui类）, 采用交叉熵损失函数和随机梯度下降优化器, 进行10轮训练，每轮包含完整训练集和测试集处理。

  3. 模型评估 ：每轮训练后在测试集上评估模型性能，计算测试损失和准确率，使用model.train()和model.eval()正确设置模型模式。

  4. 可视化与日志记录：使用TensorBoard记录训练损失、测试损失和测试准确率， 每100个训练步骤记录一次训练损失， 每轮训练后记录测试指标。

  5. 模型保存： 每轮训练后保存模型检查点， 提供了两种保存方式：完整模型和仅参数。

PS:学的似懂非懂，有问题的地方，二刷的时候，再来修改。

相关推荐
笨鸟先飞，勤能补拙3 小时前
AI 赋能网络安全：技术全景、成熟度评估与实战案例
人工智能·python·安全·web安全·网络安全·sqlite·github
长和信泰光伏储能4 小时前
京津冀光伏发电：绿色能源的未来之路
python·能源
浦信仿真大讲堂4 小时前
从重复操作到自动化闭环：如何让 CST 与 Python 真正协同起来
python·自动化·cst·仿真软件·达索软件
Gu Gu Study5 小时前
ScoutLoop开放域深度研究引擎（agent的初步设计想法）
人工智能·python
卷无止境5 小时前
写代码这件事，到底该讲究点什么？
后端·python
卷无止境6 小时前
循环复杂度到底在算什么，Python 代码怎么才能写得让人一看就懂
后端·python
lpfasd1236 小时前
MediaCrawler 项目深度分析
chrome·python·chrome devtools
Dxy12393102166 小时前
Python项目打包成EXE完整教程（PyInstaller实战避坑）
开发语言·python
bamb007 小时前
一个项目带你入门AI应用开发01
python
0566467 小时前
Python康复训练——常用标准库
开发语言·python·学习
热门推荐
01如何新建文件夹？ 电脑新建文件夹的4种方法02GitHub 镜像站点03国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）042026年7月AI圈大地震：GPT-5.6被政府限制、Claude入驻Slack、Anthropic自研芯片05微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新06AI 编程 IDE 全景解析 2026：Agent 全面接管开发链路072026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？08AI科技热点日报 | 2026年07月01日092026 年 AI 大模型 & AI 编程工具实战全总结102026上半年大模型全景技术解读：推理融合、Agent 爆发与多模态统一