卷积神经网络基础-Pytorch

PyTorch 神经网络基础

一个最简单的神经网络定义

import torch
import torch.nn as nn

class QNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super(QNetwork, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),  # 第1层:输入→隐藏层
            nn.ReLU(),                          # 激活函数(非线性)
            nn.Linear(hidden_dim, hidden_dim), # 第2层:隐藏层→隐藏层
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)  # 第3层:隐藏层→输出
        )

    def forward(self, x):
        return self.net(x)
import os                     # 导入顶级模块
import os.path as path        # 导入子模块并起别名(为了省去 os.path 的写法)
import torch
import torch.nn as nn

1. 类定义与继承

class QNetwork(nn.Module):
  • 继承自 torch.nn.Module,这是 PyTorch 中所有神经网络模块的基类
  • 通过继承,该网络能够自动管理参数、支持 GPU 迁移、实现训练/评估模式切换等

2. __init__ 方法

def __init__(self, state_dim, action_dim, hidden_dim=128):

参数

  • state_dim:状态空间的维度(输入特征数)。例如,如果状态是 4 个连续变量(CartPole 环境),则 state_dim=4
  • action_dim:动作空间的维度(输出节点数)。对于离散动作,每个动作对应一个输出 Q 值
  • hidden_dim:隐藏层的神经元数量,默认 128

网络结构

self.net = nn.Sequential(
    nn.Linear(state_dim, hidden_dim),  # 第1层:输入 → 隐藏层
    nn.ReLU(),                          # 激活函数(增加非线性)
    nn.Linear(hidden_dim, hidden_dim), # 第2层:隐藏层 → 隐藏层
    nn.ReLU(),
    nn.Linear(hidden_dim, action_dim)  # 第3层:隐藏层 → 输出层
)
  • 使用 nn.Sequential 将多个层按顺序组合,方便前向传播
  • 三层全连接网络(输入层不算在内):
    • 输入 → 第一隐藏层(state_dim → hidden_dim)+ ReLU
    • 第二隐藏层(hidden_dim → hidden_dim)+ ReLU
    • 输出层(hidden_dim → action_dim):无激活函数,因为 Q 值可以是任意实数
  • 激活函数 ReLU 引入非线性,使网络能够拟合复杂的函数关系

3. forward 方法

在 PyTorch 中,当你调用网络实例(例如 q_net(state))时,会自动执行 forward 函数。它是网络的核心计算逻辑。

def forward(self, x):
    return self.net(x)
  • 定义输入 x 如何通过网络得到输出
  • 这里直接调用 self.net(x),因为 nn.Sequential 已经包含了所有层
  • x 通常是状态张量,形状一般为 (batch_size, state_dim),输出形状为 (batch_size, action_dim),即每个动作对应的 Q 值

神经网络结构图


卷积神经网络

全连接层

nn.Linear(in_features, out_features, bias=True)
  • in_features:输入特征数
  • out_features:输出特征数
  • bias:是否使用偏置(默认 True)
layer = nn.Linear(128, 64)   # 输入128维,输出64维

卷积层

步长为1,padding=0情况下:

卷积操作示意图

具体的PyTorch实现:

import torch
import torch.nn.functional as F

input = torch.tensor([[1, 2, 0, 3, 1],
                      [0, 1, 2, 3, 1],
                      [1, 2, 1, 0, 0],
                      [5, 2, 3, 1, 1],
                      [2, 1, 0, 1, 1]])
kernel = torch.tensor([[1, 2, 1],
                       [0, 1, 0],
                       [2, 1, 0]])

input = torch.reshape(input, [1, 1, 5, 5])
kernel = torch.reshape(kernel, [1, 1, 3, 3])

output = F.conv2d(input, kernel, stride=1, padding=0)
output = torch.reshape(output, [3, 3])
print(output)

具体文档:torch.nn.functional.conv2d — PyTorch 2.12 documentation

class torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', device=None, dtype=None)

怎么理解卷积核的输出shape的计算公式

$$H_{\text{out}} = \left\lfloor
\frac{
H_{\text{in}}
+ 2 \times \text{padding}[0]
- \text{dilation}[0] \times (\text{kernel_size}[0] - 1)
- 1
}{
\text{stride}[0]
} + 1
\right\rfloor$$

滑了多少步?

输入宽度 W_in
卷积核宽度 kernel_size
每一步"吃掉" kernel_size 个格子
步与步之间间隔 stride

滑的步数 = (W_in - kernel_size) / stride

但这只是起始位置的数量,卷积核还要覆盖自身的宽度,所以输出宽度 = 步数 + 1(最后一个位置还能放一个卷积核):

输出宽度计算

可能不是整除stride,所以向下取整。

向下取整示意图

最大池化

class torch.nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)

Parameters:

  • kernel_size (int | tuple[int, int]) – the size of the window to take a max over
  • stride (int | tuple[int, int]) – the stride of the window. Default value is kernel_size
  • padding (int | tuple[int, int]) – Implicit negative infinity padding to be added on both sides
  • dilation (int | tuple[int, int]) – a parameter that controls the stride of elements in the window
  • return_indices (bool) – if True, will return the max indices along with the outputs. Useful for torch.nn.MaxUnpool2d later
  • ceil_mode (bool) – when True, will use ceil instead of floor to compute the output shape

ceil_mode操作演示

ceil_mode演示

# 最大池化
import torch
import torch.nn as nn

input = torch.tensor([[1, 2, 0, 3, 1],
                      [0, 1, 2, 3, 1],
                      [1, 2, 1, 0, 0],
                      [5, 2, 3, 1, 1],
                      [2, 1, 0, 1, 1]])

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.maxPool1 = nn.MaxPool2d(kernel_size=3, ceil_mode=True)

    def forward(self, x):
        x = self.maxPool1(x)
        return x

input = torch.reshape(input, [1, 1, 5, 5])
mynet = MyNet()
output = mynet(input)
print(output)

注意:(N, C, H, W)nn.MaxPool2d 要求的输入格式是 4维

维度 含义
N 批次
C 通道
H
W

激活函数

激活函数为神经网络引入非线性,使其能够学习和表示复杂的函数映射。

常见的激活函数:ReLU、Sigmoid、Softmax 等

现代网络 90% 的情况用 ReLU 或 GELU,其他激活函数主要出现在特定场景。

m = nn.ReLU()
input = torch.randn(2)
output = m(input)

ReLU

  • 输出范围:[0, +∞)
  • 优点:计算极快(一个 if 判断)、缓解梯度消失(正区间梯度恒为 1)
  • 缺点:Dying ReLU — 负区间永远输出 0,梯度为 0,某些神经元可能"死掉"
  • 用途:几乎所有现代网络默认首选

ReLU函数图

GELU

特点:比 ReLU 更平滑,Transformer(BERT、GPT)都用它

GELU函数图

场景 推荐
隐藏层(通用) ReLU
Transformer 类模型 GELU
二分类输出 Sigmoid
多分类输出 Softmax
怕 ReLU 死掉 LeakyReLU

线性层

线性层(Linear Layer)也叫全连接层(Fully Connected Layer)或仿射层(Affine Layer),是神经网络最基本的组成单元。

class torch.nn.Linear(in_features, out_features, bias=True, device=None, dtype=None)

参数:

  • in_features(整数)——每个输入样本的大小
  • out_features(整数)——每个输出样本的大小
  • bias(布尔值)– 如果设置为 False,则该层将不会学习加性偏置。默认值:True
m = nn.Linear(20, 30)
input = torch.randn(12, 12, 20)
output = m(input)
print(output.size())

注意: nn.Linear 总是作用于输入张量的最后一个维度。

无论输入是几维张量,nn.Linear(in, out) 只关心最后一维,其他维度全部透传(batch、seq_len 等都不会受影响)。

通常作为神经网络的尾部作为分类头:

self.classifier = nn.Sequential(
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, 10)    # 10分类
)

Dropout层

Dropout 是一种防止过拟合的正则化技术,通过在训练时随机"丢弃"部分神经元来工作。

具体做法

在训练过程中,以概率 p 随机地将输入张量的某些元素置零。对于每个前向调用,零元素都是独立选择的,并且是从伯努利分布中采样的。每次前向调用时,每个通道都会独立地被清零。

class MyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(512, 256)
        self.dropout = nn.Dropout(p=0.5)    # p=0.5,丢弃50%

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)                  # 训练时有效果,测试时自动关闭
        return x

训练 vs 测试 的区别

阶段 行为
训练 以概率 p 随机丢弃神经元,剩余的输出放大 1/(1-p)
测试 所有神经元都参与,不丢弃,但输出保持原样

为什么要放大?

因为训练时只有 (1-p) 的神经元在工作,期望输出变小了,测试时所有神经元都工作,所以要补偿回去。

Dropout放大机制

为什么能防止过拟合?

强迫网络不依赖特定神经元。

没有 Dropout 时,网络可能学到"第 5 个神经元总是负责识别猫"——这叫协同适应(co-adaptation),某个神经元坏了网络就废了。

加 Dropout 后,第 5 个神经元可能被随机丢弃,网络被迫让多个神经元共同承担识别猫的任务,结果是泛化能力更强

用在哪里

class FullNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 512),
            nn.ReLU(),
            nn.Dropout(0.3),       # 第1层后

            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.3),       # 第2层后

            nn.Linear(256, 10)
            # 输出层一般不加 Dropout
        )

通常加在隐藏层后面,不加在输出层。

场景 是否用 Dropout
数据量小,容易过拟合 ✅ 用
网络很深很宽 ✅ 用
使用 L2 正则化 可以不用
数据量很大,训练充足 ❌ 不需要
使用 BatchNorm 效果重叠,可能少用

现代实践中,BatchNorm 和 Dropout 有时会冲突,两者同时用效果可能反而变差,所以要根据自己的网络结构调一调。


损失函数

前向传播:计算损失,即实际输出与目标输出之间的差距

反向传播:为参数权重更新提供依据,梯度grad更新

MSE损失

class torch.nn.MSELoss(size_average=None, reduce=None, reduction='mean')

MSE公式

import torch
from torch.nn import MSELoss

output = torch.tensor([1, 2, 3], dtype=torch.float32)
target = torch.tensor([1, 2, 5], dtype=torch.float32)
loss_fn = MSELoss(reduction='sum')  # 两种模式:mean和sum
loss = loss_fn(output, target)      # 可以断点调试这一步后grad更新

print(loss)

交叉熵损失函数

class torch.nn.CrossEntropyLoss(weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean', label_smoothing=0.0)
import torch
from torch import optim
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear, Module, ReLU, CrossEntropyLoss
import torchvision
from torchvision.transforms import transforms

transform = transforms.Compose([
    transforms.ToTensor(),                          # 转为 Tensor,且像素值从[0,255]变成[0,1]
    transforms.Normalize((0.5, 0.5, 0.5),           # 标准化:(值 - 0.5) / 0.5
                         (0.5, 0.5, 0.5))
])
test_data = torchvision.datasets.CIFAR10(
    root='../data',           # 数据存到当前目录的 data 文件夹下
    train=False,              # 训练集(50000 张图片)
    transform=transform,      # 应用预处理
    download=True             # 没有就自动下载
)

test_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=False)

class MaoNet(Module):
    def __init__(self):
        super(MaoNet, self).__init__()
        self.model1 = Sequential(
            Conv2d(3, 32, 5, padding=2),
            ReLU(),  # 激活函数
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            ReLU(),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            ReLU(),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            ReLU(),
            Linear(64, 10)
        )

    def forward(self, x):
        x = self.model1(x)
        return x

maonet = MaoNet()
loss_fn = CrossEntropyLoss()

for data in test_loader:
    img, label = data
    output = maonet(img)
    loss = loss_fn(output, label)
    print(loss)  # 打印损失

优化器

在深度学习和机器学习中,优化器(Optimizer) 是用于更新模型参数(如权重和偏置)以最小化损失函数的算法。

torch.optim 是一个实现了多种优化算法的软件包。

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)

用于计算梯度之后,调用 optimizer.step(),进行参数更新。

通常是这样:

for input, target in dataset:
    optimizer.zero_grad()        # 1. 梯度重置为0
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()              # 2. 梯度更新
    optimizer.step()             # 3. 根据梯度更新参数

优化器工作流程

优化器对比


网络架构搭建

CIFAR10 Net

CIFAR10网络结构

import torch
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear, Module, ReLU

class MaoNet(Module):
    def __init__(self):
        super(MaoNet, self).__init__()
        self.model1 = Sequential(
            Conv2d(3, 32, 5, padding=2),
            ReLU(),  # 激活函数
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            ReLU(),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            ReLU(),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            ReLU(),
            Linear(64, 10)
        )

    def forward(self, x):
        x = self.model1(x)
        return x

maonet = MaoNet()
input = torch.ones([64, 3, 32, 32])
output = maonet(input)
print(output.shape)

VGG 16 Net

VGG16网络结构

import torch
from torch.nn import Module, Sequential, Linear, ReLU, Conv2d, MaxPool2d, Softmax, Flatten

class vgg16_me(Module):
    def __init__(self):
        super(vgg16_me, self).__init__()
        self.net = Sequential(
            # Block1
            Conv2d(3, 64, 3, padding=1),
            ReLU(),
            Conv2d(64, 64, 3, padding=1),
            ReLU(),
            MaxPool2d(2),

            # Block2
            Conv2d(64, 128, 3, padding=1),
            ReLU(),
            Conv2d(128, 128, 3, padding=1),
            ReLU(),
            MaxPool2d(2),

            # Block3
            Conv2d(128, 256, 3, padding=1),
            ReLU(),
            Conv2d(256, 256, 3, padding=1),
            ReLU(),
            Conv2d(256, 256, 3, padding=1),
            ReLU(),
            MaxPool2d(2),

            # Block4
            Conv2d(256, 512, 3, padding=1),
            ReLU(),
            Conv2d(512, 512, 3, padding=1),
            ReLU(),
            Conv2d(512, 512, 3, padding=1),
            ReLU(),
            MaxPool2d(2),

            # Block5
            Conv2d(512, 512, 3, padding=1),
            ReLU(),
            Conv2d(512, 512, 3, padding=1),
            ReLU(),
            Conv2d(512, 512, 3, padding=1),
            ReLU(),
            MaxPool2d(2),
        )
        self.classify = Sequential(
            Flatten(),
            Linear(25088, 4096), ReLU(),
            Linear(4096, 4096), ReLU(),
            Linear(4096, 1000),
        )
        # 如果推理时需要 Softmax,可以手动添加
        self.softmax = Softmax(dim=1)

    def forward(self, x):
        x = self.net(x)
        x = self.classify(x)
        return x

vgg = vgg16_me()
input = torch.zeros([64, 3, 224, 224])
output = vgg(input)
print(output.shape)

ResNet 残差网络

出现原因: 网络结构越深,效果反而更差

为什么? 深度网络有梯度消失的问题

损失 → 第56层 → 第55层 → ... → 第2层 → 第1层

损失函数从后往前由复合函数求偏导的乘法法则,梯度从后往前传,每经过一层就要乘一个数。乘的数如果小于1:

0.9 × 0.9 × 0.9 × ...(乘了56次)→ 几乎为0

第一层根本收不到有效的梯度信号,学不动

解决方法(ResNet思路): 与其让网络直接学输出,不如让它学"差值"(残差)。

  • 普通网络:输出 = F(x) ← 网络要直接学 F(x)
  • ResNet:输出 = x + F(x) ← 网络只需要学 F(x) = 输出 - x

ResNet残差块

为什么学差值就简单了?

  • 恒等映射对普通网络很难,对 ResNet 只需要把权重学成0

梯度怎么流?

  • 普通网络:梯度 = ∂F/∂x,乘多层可能趋近于0
  • ResNet:梯度 = 1 + ∂F/∂x永远有一条"高速公路"让梯度直接流回去

ResNet梯度流

ResNet结构

import torch
from torch.nn import Module, Sequential, Linear, ReLU, Conv2d, MaxPool2d, AdaptiveAvgPool2d, Flatten

class Bottleneck(Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super(Bottleneck, self).__init__()
        mid_channels = out_channels // 4

        self.core = Sequential(
            Conv2d(in_channels, mid_channels, kernel_size=1, bias=False),
            ReLU(),
            Conv2d(mid_channels, mid_channels, kernel_size=3, stride=stride, padding=1, bias=False),
            ReLU(),
            Conv2d(mid_channels, out_channels, kernel_size=1, bias=False)
        )

        # 残差连接:当输入输出维度不同时,用1x1卷积把x映射到相同维度
        self.shortcut = Sequential()
        if in_channels != out_channels:
            self.shortcut = Sequential(
                Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)
            )

        self.relu = ReLU()

    def forward(self, x):
        out = self.core(x)
        residual = self.shortcut(x)   # 维度对齐
        out = self.relu(out + residual)
        return out

class ResNet50(Module):
    def __init__(self, num_classes=1000):
        super(ResNet50, self).__init__()

        # ── header: 224×224 → 56×56 ──
        self.header = Sequential(
            Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
            ReLU(),
            MaxPool2d(kernel_size=3, stride=2, padding=1)
        )

        # ── body: 4个stage,分别有 3, 4, 6, 3 个 Bottleneck ──
        self.stage1 = self._make_stage(64, 256, num_blocks=3, stride=1)
        self.stage2 = self._make_stage(256, 512, num_blocks=4, stride=2)
        self.stage3 = self._make_stage(512, 1024, num_blocks=6, stride=2)
        self.stage4 = self._make_stage(1024, 2048, num_blocks=3, stride=2)

        # ── tail: 全局平均池化 + 全连接 ──
        self.classify = Sequential(
            AdaptiveAvgPool2d(1),       # 2048 × 7 × 7 → 2048 × 1 × 1
            Flatten(),                  # → 2048
            Linear(2048, num_classes)    # → 1000
        )

    def _make_stage(self, in_channels, out_channels, num_blocks, stride):
        layers = []
        layers.append(Bottleneck(in_channels, out_channels, stride))
        for _ in range(1, num_blocks):
            layers.append(Bottleneck(out_channels, out_channels, stride=1))
        return Sequential(*layers)

    def forward(self, x):
        x = self.header(x)       # 224 → 56
        x = self.stage1(x)       # 56 × 56
        x = self.stage2(x)       # 28 × 28
        x = self.stage3(x)       # 14 × 14
        x = self.stage4(x)       # 7 × 7
        x = self.classify(x)     # 1000
        return x

# 测试
model = ResNet50(num_classes=1000)
input = torch.zeros([64, 3, 224, 224])
output = model(input)
print(output.shape)   # torch.Size([64, 1000])

网络模型的使用

torchvision提供一系列的视觉相关的模型和模型的预训练参数

Models and pre-trained weights — Torchvision 0.27 documentation

实例化预训练模型会将权重下载到缓存目录。可以使用 TORCH_HOME 环境变量设置此目录。

import os
from torchvision.models import vgg16

# 设置模型保存路径
os.environ['TORCH_HOME'] = '../weights'

# deprecated
# model_false = vgg16(pretrained=False)
# model_true = vgg16(pretrained=True)

model_false = vgg16(weights=None)
model_true = vgg16(weights='DEFAULT', progress=True)
print(model_true)

初始化预训练模型

加载权重文件:

# 1. 先创建空模型
model_true = vgg16(weights=None)   # 不下载,不加载预训练权重

# 2. 加载本地权重文件
model.load_state_dict(torch.load('model.pth', weights_only=True))  # .pth 文件路径

strict=False 让你只加载能对上的部分,对不上的保持随机初始化,然后用新数据微调。这就是迁移学习的核心思路。

model_state_dict = torch.load('../weights/vgg16-397923af.pth', weights_only=True)
model_true.load_state_dict(model_state_dict, strict=False)

如何修改已有的模型

以VGG16为例:

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

尾部增加模块,使用 add_module

model.add_module("classify head", Linear(1000, 10))

# 或者增加到分类头里面去
model.classifier.add_module("classify head", Linear(1000, 10))

修改中间层,直接使用数组下标修改:

model.classifier[6] = Linear(4096, 10)

模型保存

模型保存有两种方式:

torch.save(model, "./XXX.pth")
torch.save(model.state_dict(), "./XXX.pth")  # 官方推荐
torch.save(model, ...) torch.save(model.state_dict(), ...)
保存的是什么 整个模型(结构 + 权重) 只有权重参数
文件大小 大(包含模型类定义) 小(只有数据)
加载方式 torch.load(path) 直接用 需要先建模型再填权重

训练部分层

方法一:冻结部分层

# 先冻结所有参数
for param in model.parameters():
    param.requires_grad = False

# 只解冻新的全连接层
for param in model.fc.parameters():
    param.requires_grad = True

方法二:直接将需要训练的层传入优化器

optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)
# 或者用 Adam
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)

模型训练总流程

import torch
import torchvision
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader
from torchvision import transforms
from model import MaoNet

# 1. 预处理
transform = transforms.Compose([
    transforms.ToTensor(),                          # 转为 Tensor,且像素值从[0,255]变成[0,1]
    transforms.Normalize((0.5, 0.5, 0.5),           # 标准化:(值 - 0.5) / 0.5
                         (0.5, 0.5, 0.5))
])

# 2. 数据集
train_data = torchvision.datasets.CIFAR10(
    root='./data',           # 数据存到当前目录的 data 文件夹下
    train=True,              # 训练集(50000 张图片)
    transform=transform,     # 应用预处理
    download=True            # 没有就自动下载
)

test_data = torchvision.datasets.CIFAR10(
    root='./data',
    train=False,             # 测试集(10000 张图片)
    transform=transform,
    download=True
)

# 3. DataLoader:把数据分批次加载
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader  = DataLoader(test_data,  batch_size=64, shuffle=False)

# CIFAR-10 的 10 个类别
classes = ('飞机', '汽车', '鸟', '猫', '鹿', '狗', '青蛙', '马', '船', '卡车')

# ============================================================
# 4. 初始化模型、损失函数、优化器
# ============================================================
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

maonet = MaoNet().to(device)
loss_fn = CrossEntropyLoss()                       # 分类用交叉熵损失
optimizer = torch.optim.Adam(maonet.parameters(), lr=0.001)  # Adam 优化器

# ============================================================
# 5. 训练
# ============================================================
total_train_step = 0

for epoch in range(10):  # 训练 10 轮
    maonet.train()  # 切换到训练模式(Dropout/BatchNorm 会不同)
    total_loss = 0

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        # 前向传播
        outputs = maonet(images)
        loss = loss_fn(outputs, labels)

        # 反向传播
        optimizer.zero_grad()  # 5.1 清空之前的梯度
        loss.backward()        # 5.2 计算梯度
        optimizer.step()       # 5.3 更新参数

        total_loss += loss.item()
        total_train_step += 1

    avg_loss = total_loss / len(train_loader)

    # ============================================================
    # 6. 每轮结束后在测试集上验证
    # ============================================================
    maonet.eval()  # 6.1 切换到测试模式
    total_correct = 0
    total_samples = 0

    with torch.no_grad():  # 6.2 测试时不需要计算梯度
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = maonet(images)
            _, predicted = torch.max(outputs, 1)  # 取最大值对应的类别

            total_correct += (predicted == labels).sum().item()
            total_samples += labels.size(0)

    accuracy = total_correct / total_samples * 100
    print(f"Epoch [{epoch+1}/10]  Loss: {avg_loss:.4f}  测试准确率: {accuracy:.2f}%")

print("训练完成!")

整体框架伪代码

# ============================================================
# 1. 准备数据
# ============================================================
transform = transforms.Compose([...])                              # 数据预处理
train_dataset = Dataset(root='...', train=True, transform=transform)  # 训练集
test_dataset  = Dataset(root='...', train=False, transform=transform) # 测试集
train_loader  = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader   = DataLoader(test_dataset,  batch_size=64, shuffle=False)

# ============================================================
# 2. 定义模型
# ============================================================
class MyNet(nn.Module):
    def __init__(self):
        # 定义网络结构:Conv → ReLU → Pool → ... → FC → Output

    def forward(self, x):
        # 前向传播,定义数据怎么流过网络
        return x

# ============================================================
# 3. 初始化三件套
# ============================================================
model     = MyNet()
loss_fn   = nn.CrossEntropyLoss()                    # 损失函数
optimizer = torch.optim.Adam(model.parameters())      # 优化器

# ============================================================
# 4. 训练循环
# ============================================================
for epoch in range(N):
    model.train()                     # 切训练模式
    for images, labels in train_loader:
        outputs = model(images)        # 前向传播 → 自动建计算图
        loss = loss_fn(outputs, labels) # 算损失

        optimizer.zero_grad()          # 清梯度
        loss.backward()                # 算梯度(沿计算图反向传播)
        optimizer.step()               # 更新参数

    # ============================================================
    # 5. 测试
    # ============================================================
    model.eval()                       # 切测试模式
    with torch.no_grad():             # 不需要梯度
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            # 统计正确率

训练流程:

epoch 外循环
  ↓
  batch 内循环
    ↓
    前向传播 → 损失 → 清梯度 → 反向传播 → 更新参数
posted @ 2026-06-11 20:47  MaoShine  阅读(15)  评论(0)    收藏  举报