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"))

相关推荐
傲笑风6 小时前
【openvino】tinybert基于openvino服务化部署(四)
人工智能·python·自然语言处理·nlp·bert·openvino
维基框架7 小时前
WIKI 知识库 v1.1.1 正式发布
人工智能·python
谢白羽11 小时前
SGLang的AWQ量化笔记
笔记·python·sglang
迷迭香yy11 小时前
基金档案数据工程实战从收入分析到持仓穿透的Python解析 IG50免费开源股票数据API接口
开发语言·python
清水白石00811 小时前
Python 类型设计深度解析:TypedDict 能否替代 dataclass?从 JSON 数据边界到 API 设计的最佳实践
java·python·json
️学习的小王12 小时前
Git项目提交忽略文件怎么做?以Python项目为例,详解.gitignore
git·python·elasticsearch
前端 贾公子12 小时前
第09章:上下文与记忆 (6)
开发语言·前端·python
evans在进步12 小时前
Java 常用设计模式入门:建造者、工厂、单例、外观与代理
java·python·设计模式
jayson.h13 小时前
PDF 合并+添加页码 相关库、类、函数
开发语言·前端·python
典典分享指南13 小时前
用飞书多维表格搭建内容管理看板
java·python·r语言·composer·symfony