PyTorch零基础】模型修改、添加网络层、两种模型保存与加载方式

前言

在深度学习训练中,我们经常需要:

  • 修改官方预训练模型(VGG、ResNet)

  • 给网络新增网络层

  • 保存训练好的模型、断点续训

  • 加载模型继续训练/推理

PyTorch 提供了两种保存模型、两种加载模型的方式,新手极易混淆、极易报错。

本文结合 VGG16 实战代码,一次性彻底讲透,以后永远不用死记硬背。

一、基础认知:PyTorch 模型构成

一个完整的神经网络包含两部分:

  1. 模型结构:卷积、池化、全连接、Sequential 等网络骨架

  2. 模型参数(权重):weight、bias(可训练参数)

两种保存方式本质区别:

  • 方式1:保存 结构 + 权重

  • 方式2:只保存 权重参数(官方推荐)

二、实战1:加载官方模型 + 自定义修改网络

1. 加载原生VGG16(无预训练权重)

import torch from torchvision import models # 只加载网络结构,权重随机初始化 vgg16 = models.vgg16(weights=None) print(vgg16)

此时 VGG16 默认输出是 1000 分类(ImageNet数据集)。

2. 给模型【新增网络层】(重点)

如果我要改成 10分类(CIFAR10),可以直接给模型末尾加一层:

# 给模型动态添加一层全连接 vgg16.add_module("add_linear", torch.nn.Linear(1000, 10)) print(vgg16)

add_module(层名, 层结构)

执行后,模型结构多了:add_linear: Linear(1000, 10)

✅ 此时模型结构被永久修改,权重也多出两组参数:

  • add_linear.weight

  • add_linear.bias

👉 后续保存权重,就会包含这两个参数,加载时结构必须一致,否则报错!

三、PyTorch 两种模型保存方式(核心重点)

方式一:保存整个模型对象(结构+权重)

# 保存整个模型 torch.save(vgg16, "vgg16_method1.pth")

保存内容:网络结构 + 所有权重参数

优点:加载不用写模型结构,一行直接用

致命缺点

  • 兼容性极差,PyTorch版本更新就炸(2.6版本后严格限制)

  • 存在安全风险,新版本默认禁止加载

工作、比赛、项目一律不推荐

方式二:只保存模型参数(state_dict)【官方推荐】

# 只保存权重参数字典 torch.save(vgg16.state_dict(), "vgg16_method2.pth")

保存内容:仅保存所有 weight、bias(字典形式)

优点

  • 体积最小、纯参数、无版本兼容问题

  • 绝对安全,工业界通用标准

缺点 :加载前必须手动搭建一模一样的模型结构

唯一推荐的保存方式

四、两种对应的加载方式(必须一一对应)

方式1:加载【完整模型文件】(对应保存方式1)

import torch # 新版PyTorch必须加 weights_only=False model = torch.load("vgg16_method1.pth", weights_only=False) print(model)

直接加载出完整模型,不需要搭建网络。

⚠️ 仅限自己本地测试使用!

方式2:加载【权重参数文件】(对应保存方式2)

最容易报错、最重要!!!

规则:加载时的模型结构 = 保存时的模型结构(必须完全一致)

因为我保存时多加了一层 add_linear,所以加载时也要加!

import torch from torchvision import models # 1. 先搭建【和保存时一模一样】的结构 model = models.vgg16(weights=None) model.add_module("add_linear", torch.nn.Linear(1000, 10)) # 2. 再加载权重 model.load_state_dict(torch.load("vgg16_method2.pth")) print(model)

五、你之前报错的终极解释(必看)

报错内容

Unexpected key(s) in state_dict: "add_linear.weight", "add_linear.bias"

报错原因(100%精准)

  • 保存时:模型多了 add_linear 层,权重文件里存了这两个参数

  • 加载时:你只创建了原生 VGG16,没有 add_linear

  • 结构对不上,权重参数无处安放 → 直接报错

解决方法:加载前,补全所有自定义层!

六、拓展:常用两种模型修改方式

方式1:add_module 新增层(你学的)

适合:模型末尾追加分类头

方式2:直接替换原有层(更常用)

# 直接替换VGG最后一层分类器 vgg16.classifier[6] = torch.nn.Linear(4096, 10)

七、最终总结(可背诵)

1. 两种保存区别

  • torch.save(model, path):保存结构+权重,省事但不通用、新版本易报错

  • torch.save(model.state_dict(), path):只保存权重,官方标准、工业唯一推荐

2. 两种加载规则

  • 完整模型文件 → torch.load() 直接加载

  • 权重参数文件 → 先搭结构、后载权重

3. 报错核心口诀

权重里有的层,模型里必须有;权重里没有的层,模型不能乱加。

4. 工作最佳实践

永远只用:state_dict 保存 + 先建结构再加载

八、完整可运行模板(以后直接复制)

保存模板

model = models.vgg16(weights=None) model.add_module("add_linear", torch.nn.Linear(1000,10)) torch.save(model.state_dict(), "best.pth")

加载模板

model = models.vgg16(weights=None) model.add_module("add_linear", torch.nn.Linear(1000,10)) model.load_state_dict(torch.load("best.pth"))

相关推荐
IT小盘43 分钟前
03-大模型API不只是发送Prompt-流式输出超时重试与异常处理
人工智能·python·prompt
Zzz不能停2 小时前
个人博客系统系统---测试报告【笔耕云】
python·功能测试·自动化·压力测试
互联网中的一颗神经元3 小时前
小白python入门 - 39. 采集流水线小项目
开发语言·python
郭老二4 小时前
【Python】常用模块:xmlrpc
python
名字还没想好☜4 小时前
Python itertools 实战:用 groupby、chain、islice 优雅处理大数据流
linux·windows·python·迭代器·itertools
威联通网络存储5 小时前
TS-h2490FU在面板制造Array段AOI缺陷画廊中的并联
python·制造
接针6 小时前
UV 常用命令
python·uv
无垠的广袤7 小时前
【工业树莓派 CM0 Dev Board】扩展板设计
linux·python·嵌入式硬件·pcb设计·模块化·传感器
Yolanda_20228 小时前
在vscode终端中可使用pip,但是cmd终端中找不到命令 pip问题的解决
vscode·python·conda·pip
2301_764441338 小时前
用动力学系统(微分方程)为 Kernberg 的客体关系单元提供数学化的操作定义,把“自体—客体“这对心理结构建模成一个二维耦合系统
数据结构·python·算法·数学建模