PyTorch预训练模型加载全攻略
PyTorch加载预训练模型是深度学习中的常见操作,本文将详细介绍在PyTorch中加载预训练模型的各种方法,助力开发者快速上手。针对ResNet、VGG等常见模型,推荐使用torchvision.models直接加载预训练版本,并通过weights参数指定权重类型,方便快捷。对于自定义模型,则需手动加载权重,在创建模型后使用load_state_dict函数,并设置strict=False以忽略不匹配部分,灵活应对模型结构的差异。此外,文章还将讲解微调模型时如何冻结部分层,以及推荐使用模型状态字典进行保存和加载,以提高灵活性并避免设备不一致问题。掌握这些关键点,能够有效解决PyTorch中加载预训练模型的各种问题,提升开发效率。
在PyTorch中加载预训练模型的关键在于理解模型结构与权重的匹配关系。1. 对于ResNet、VGG等常见模型,可通过torchvision.models直接加载预训练版本,使用weights参数指定权重类型更推荐;2. 自定义模型需手动加载权重,创建模型后使用load_state_dict并设置strict=False忽略不匹配部分;3. 微调时可冻结部分层,通过设置requires_grad=False锁定参数;4. 推荐仅保存和加载模型状态字典以提高灵活性,并注意设备一致性问题。掌握这些要点即可应对大多数情况。
在PyTorch中加载预训练模型其实并不复杂,关键在于理解模型结构和权重之间的关系。如果你用的是常见的模型架构(比如ResNet、VGG等),通常可以直接从torchvision.models
中获取对应的预训练版本。而对于自定义模型或部分修改后的模型,就需要更细致地处理权重加载方式。

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

例如:
import torchvision.models as models model = models.resnet18(pretrained=True)
这样一行代码就能加载一个带预训练权重的ResNet18模型。如果你想换成其他变种,比如resnet50
或者vgg16
,只需要替换名字即可。

需要注意的是:从 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学习网公众号,给大家分享更多科技周边知识!

- 上一篇
- Promise封装XMLHttpRequest方法详解

- 下一篇
- 豆包生成Python数据库连接代码方法详解
-
- 科技周边 · 人工智能 | 2小时前 |
- Deepseek满血版+AIDungeon,开启故事新纪元
- 170浏览 收藏
-
- 科技周边 · 人工智能 | 2小时前 |
- 豆包AI编程工具入门指南
- 458浏览 收藏
-
- 科技周边 · 人工智能 | 2小时前 |
- Deepseek满血版搭配Prezi,打造惊艳演示
- 444浏览 收藏
-
- 科技周边 · 人工智能 | 2小时前 |
- AI证件照换背景的5种技巧
- 420浏览 收藏
-
- 科技周边 · 人工智能 | 2小时前 |
- 通义千问API接入与对话测试指南
- 175浏览 收藏
-
- 科技周边 · 人工智能 | 3小时前 |
- 多模态AI如何解析化学式与公式理解技术
- 477浏览 收藏
-
- 科技周边 · 人工智能 | 3小时前 | 一维金属纳米结构 金纳米棒 相干振动 光学性质 飞秒时间分辨瞬态吸收光谱
- 中国科大一维金属纳米研究获突破
- 215浏览 收藏
-
- 科技周边 · 人工智能 | 3小时前 |
- 豆包AI风格切换全攻略
- 309浏览 收藏
-
- 前端进阶之JavaScript设计模式
- 设计模式是开发人员在软件开发过程中面临一般问题时的解决方案,代表了最佳的实践。本课程的主打内容包括JS常见设计模式以及具体应用场景,打造一站式知识长龙服务,适合有JS基础的同学学习。
- 542次学习
-
- GO语言核心编程课程
- 本课程采用真实案例,全面具体可落地,从理论到实践,一步一步将GO核心编程技术、编程思想、底层实现融会贯通,使学习者贴近时代脉搏,做IT互联网时代的弄潮儿。
- 510次学习
-
- 简单聊聊mysql8与网络通信
- 如有问题加微信:Le-studyg;在课程中,我们将首先介绍MySQL8的新特性,包括性能优化、安全增强、新数据类型等,帮助学生快速熟悉MySQL8的最新功能。接着,我们将深入解析MySQL的网络通信机制,包括协议、连接管理、数据传输等,让
- 498次学习
-
- JavaScript正则表达式基础与实战
- 在任何一门编程语言中,正则表达式,都是一项重要的知识,它提供了高效的字符串匹配与捕获机制,可以极大的简化程序设计。
- 487次学习
-
- 从零制作响应式网站—Grid布局
- 本系列教程将展示从零制作一个假想的网络科技公司官网,分为导航,轮播,关于我们,成功案例,服务流程,团队介绍,数据部分,公司动态,底部信息等内容区块。网站整体采用CSSGrid布局,支持响应式,有流畅过渡和展现动画。
- 484次学习
-
- 边界AI平台
- 探索AI边界平台,领先的智能AI对话、写作与画图生成工具。高效便捷,满足多样化需求。立即体验!
- 397次使用
-
- 免费AI认证证书
- 科大讯飞AI大学堂推出免费大模型工程师认证,助力您掌握AI技能,提升职场竞争力。体系化学习,实战项目,权威认证,助您成为企业级大模型应用人才。
- 405次使用
-
- 茅茅虫AIGC检测
- 茅茅虫AIGC检测,湖南茅茅虫科技有限公司倾力打造,运用NLP技术精准识别AI生成文本,提供论文、专著等学术文本的AIGC检测服务。支持多种格式,生成可视化报告,保障您的学术诚信和内容质量。
- 543次使用
-
- 赛林匹克平台(Challympics)
- 探索赛林匹克平台Challympics,一个聚焦人工智能、算力算法、量子计算等前沿技术的赛事聚合平台。连接产学研用,助力科技创新与产业升级。
- 641次使用
-
- 笔格AIPPT
- SEO 笔格AIPPT是135编辑器推出的AI智能PPT制作平台,依托DeepSeek大模型,实现智能大纲生成、一键PPT生成、AI文字优化、图像生成等功能。免费试用,提升PPT制作效率,适用于商务演示、教育培训等多种场景。
- 549次使用
-
- GPT-4王者加冕!读图做题性能炸天,凭自己就能考上斯坦福
- 2023-04-25 501浏览
-
- 单块V100训练模型提速72倍!尤洋团队新成果获AAAI 2023杰出论文奖
- 2023-04-24 501浏览
-
- ChatGPT 真的会接管世界吗?
- 2023-04-13 501浏览
-
- VR的终极形态是「假眼」?Neuralink前联合创始人掏出新产品:科学之眼!
- 2023-04-30 501浏览
-
- 实现实时制造可视性优势有哪些?
- 2023-04-15 501浏览