当前位置:首页 > 文章列表 > 科技周边 > 人工智能 > PyTorch预训练模型加载全攻略

PyTorch预训练模型加载全攻略

2025-07-22 23:53:19 0浏览 收藏

最近发现不少小伙伴都对科技周边很感兴趣,所以今天继续给大家介绍科技周边相关的知识,本文《PyTorch加载预训练模型方法详解》主要内容涉及到等等知识点,希望能帮到你!当然如果阅读本文时存在不同想法,可以在评论中表达,但是请勿使用过激的措辞~

在PyTorch中加载预训练模型的关键在于理解模型结构与权重的匹配关系。1. 对于ResNet、VGG等常见模型,可通过torchvision.models直接加载预训练版本,使用weights参数指定权重类型更推荐;2. 自定义模型需手动加载权重,创建模型后使用load_state_dict并设置strict=False忽略不匹配部分;3. 微调时可冻结部分层,通过设置requires_grad=False锁定参数;4. 推荐仅保存和加载模型状态字典以提高灵活性,并注意设备一致性问题。掌握这些要点即可应对大多数情况。

如何在PyTorch中加载预训练模型 PyTorch模型加载与微调实践

在PyTorch中加载预训练模型其实并不复杂,关键在于理解模型结构和权重之间的关系。如果你用的是常见的模型架构(比如ResNet、VGG等),通常可以直接从torchvision.models中获取对应的预训练版本。而对于自定义模型或部分修改后的模型,就需要更细致地处理权重加载方式。

如何在PyTorch中加载预训练模型 PyTorch模型加载与微调实践

使用torchvision快速加载常见预训练模型

对于像ResNet、AlexNet这类经典模型,推荐使用torchvision.models中的接口。这些模型默认提供是否加载预训练权重的选项。

如何在PyTorch中加载预训练模型 PyTorch模型加载与微调实践

例如:

import torchvision.models as models

model = models.resnet18(pretrained=True)

这样一行代码就能加载一个带预训练权重的ResNet18模型。如果你想换成其他变种,比如resnet50或者vgg16,只需要替换名字即可。

如何在PyTorch中加载预训练模型 PyTorch模型加载与微调实践

需要注意的是:从 PyTorch 1.8 开始,pretrained参数被逐步弃用,取而代之的是通过weights参数指定具体权重类型。比如:

model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)

这种方式更清晰,也方便你选择不同来源的预训练权重。

自定义模型加载预训练权重

如果你对模型结构做了改动,比如改了最后一层分类器的输出维度,就不能直接使用pretrained=True,而是需要手动加载权重。

一般做法是先创建你的模型结构,然后加载原始权重文件,并忽略不匹配的部分。

步骤如下:

  • 创建模型结构
  • 加载预训练权重字典
  • 使用load_state_dict并设置strict=False

示例代码:

model = MyModifiedResNet()
pretrained_dict = torch.load('resnet18.pth')
model.load_state_dict(pretrained_dict, strict=False)

这里的关键是确保大部分层的名字和结构都能对应上,否则会报错。如果模型结构差异太大,可能需要手动映射某些层。

微调时冻结部分层

微调模型时,有时候我们只想训练最后几层,而保持前面的特征提取层不变。这时候可以通过设置requires_grad来冻结部分参数。

以ResNet为例,假设你想冻结所有层,只训练最后的全连接层:

for param in model.parameters():
    param.requires_grad = False

# 只训练最后的fc层
model.fc = nn.Linear(512, new_num_classes)

这样,在反向传播过程中,只有model.fc的参数会被更新。

当然,也可以有选择地冻结某些层,而不是全部。比如只冻结前几层卷积层,保留后面层可训练。

权重保存与加载格式注意事项

PyTorch模型保存的常见方式有两种:

  • 保存整个模型:torch.save(model, PATH)
  • 仅保存模型状态字典:torch.save(model.state_dict(), PATH)

推荐使用第二种方式,因为它更灵活,特别是在模型类定义发生变化后仍然能加载权重。

加载时也要注意设备问题。如果保存的是GPU上的模型,而在CPU环境下加载,需要加上map_location='cpu'

model.load_state_dict(torch.load(PATH, map_location='cpu'))

否则可能会遇到设备不匹配的问题。


基本上就这些操作了。加载预训练模型本身不难,但要注意结构匹配、设备一致性和是否需要冻结部分层。掌握这几个点,大多数情况都能应对。

好了,本文到此结束,带大家了解了《PyTorch预训练模型加载全攻略》,希望本文对你有所帮助!关注golang学习网公众号,给大家分享更多科技周边知识!

GinEchoBeego框架对比评测GinEchoBeego框架对比评测
上一篇
GinEchoBeego框架对比评测
PythonElementTree解析教程详解
下一篇
PythonElementTree解析教程详解
查看更多
最新文章
查看更多
课程推荐
  • 前端进阶之JavaScript设计模式
    前端进阶之JavaScript设计模式
    设计模式是开发人员在软件开发过程中面临一般问题时的解决方案,代表了最佳的实践。本课程的主打内容包括JS常见设计模式以及具体应用场景,打造一站式知识长龙服务,适合有JS基础的同学学习。
    542次学习
  • GO语言核心编程课程
    GO语言核心编程课程
    本课程采用真实案例,全面具体可落地,从理论到实践,一步一步将GO核心编程技术、编程思想、底层实现融会贯通,使学习者贴近时代脉搏,做IT互联网时代的弄潮儿。
    511次学习
  • 简单聊聊mysql8与网络通信
    简单聊聊mysql8与网络通信
    如有问题加微信:Le-studyg;在课程中,我们将首先介绍MySQL8的新特性,包括性能优化、安全增强、新数据类型等,帮助学生快速熟悉MySQL8的最新功能。接着,我们将深入解析MySQL的网络通信机制,包括协议、连接管理、数据传输等,让
    498次学习
  • JavaScript正则表达式基础与实战
    JavaScript正则表达式基础与实战
    在任何一门编程语言中,正则表达式,都是一项重要的知识,它提供了高效的字符串匹配与捕获机制,可以极大的简化程序设计。
    487次学习
  • 从零制作响应式网站—Grid布局
    从零制作响应式网站—Grid布局
    本系列教程将展示从零制作一个假想的网络科技公司官网,分为导航,轮播,关于我们,成功案例,服务流程,团队介绍,数据部分,公司动态,底部信息等内容区块。网站整体采用CSSGrid布局,支持响应式,有流畅过渡和展现动画。
    484次学习
查看更多
AI推荐
  • AI歌曲生成器:免费在线创作,一键生成原创音乐
    AI歌曲生成器
    AI歌曲生成器,免费在线创作,简单模式快速生成,自定义模式精细控制,多种音乐风格可选,免版税商用,让您轻松创作专属音乐。
    20次使用
  • MeloHunt:免费AI音乐生成器,零基础创作高品质音乐
    MeloHunt
    MeloHunt是一款强大的免费在线AI音乐生成平台,让您轻松创作原创、高质量的音乐作品。无需专业知识,满足内容创作、影视制作、游戏开发等多种需求。
    20次使用
  • 满分语法:免费在线英语语法检查器 | 论文作文邮件一键纠错润色
    满分语法
    满分语法是一款免费在线英语语法检查器,助您一键纠正所有英语语法、拼写、标点错误及病句。支持论文、作文、翻译、邮件语法检查与文本润色,并提供详细语法讲解,是英语学习与使用者必备工具。
    31次使用
  • 易销AI:跨境电商AI营销专家 | 高效文案生成,敏感词规避,多语言覆盖
    易销AI-专为跨境
    易销AI是专为跨境电商打造的AI营销神器,提供多语言广告/产品文案高效生成、精准敏感词规避,并配备定制AI角色,助力卖家提升全球市场广告投放效果与回报率。
    31次使用
  • WisFile:免费AI本地文件批量重命名与智能归档工具
    WisFile-批量改名
    WisFile是一款免费AI本地工具,专为解决文件命名混乱、归类无序难题。智能识别关键词,AI批量重命名,100%隐私保护,让您的文件井井有条,触手可及。
    31次使用
微信登录更方便
  • 密码登录
  • 注册账号
登录即同意 用户协议隐私政策
返回登录
  • 重置密码