PyTorch 的预训练,是时候学习一下了

前言

最近使用 PyTorch 感觉妙不可言,有种当初使用 Keras 的快感,而且速度还不慢。各种设计直接简洁,方便研究,比 tensorflow 的臃肿好多了。今天让我们来谈谈 PyTorch 的预训练,主要是自己写代码的经验以及论坛 PyTorch Forums 上的一些回答的总结整理。

直接加载预训练模型

如果我们使用的模型和原模型完全一样,那么我们可以直接加载别人训练好的模型:

my_resnet = MyResNet(*args, **kwargs) my_resnet.load_state_dict(torch.load("my_resnet.pth"))

当然这样的加载方法是基于 PyTorch 推荐的存储模型的方法:

torch.save(my_resnet.state_dict(), "my_resnet.pth")

还有第二种加载方法:

my_resnet = torch.load("my_resnet.pth")

加载部分预训练模型

其实大多数时候我们需要根据我们的任务调节我们的模型,所以很难保证模型和公开的模型完全一样,但是预训练模型的参数确实有助于提高训练的准确率,为了结合二者的优点,就需要我们加载部分预训练模型。

pretrained_dict = model_zoo.load_url(model_urls['resnet152']) model_dict = model.state_dict() # 将 pretrained_dict 里不属于 model_dict 的键剔除掉 pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} # 更新现有的 model_dict model_dict.update(pretrained_dict) # 加载我们真正需要的 state_dict model.load_state_dict(model_dict)

因为需要剔除原模型中不匹配的键,也就是层的名字,所以我们的新模型改变了的层需要和原模型对应层的名字不一样,比如:resnet 最后一层的名字是 fc(PyTorch 中),那么我们修改过的 resnet 的最后一层就不能取这个名字,可以叫 fc_

微改基础模型预训练

对于改动比较大的模型,我们可能需要自己实现一下再加载别人的预训练参数。但是,对于一些基本模型 PyTorch 中已经有了,而且我只想进行一些小的改动那么怎么办呢?难道我又去实现一遍吗?当然不是。

我们首先看看怎么进行微改模型。

微改基础模型

PyTorch 中的 torchvision 里已经有很多常用的模型了,可以直接调用:

  • AlexNet
  • VGG
  • ResNet
  • SqueezeNet
  • DenseNet

import torchvision.models as models resnet18 = models.resnet18() alexnet = models.alexnet() squeezenet = models.squeezenet1_0() densenet = models.densenet_161()

但是对于我们的任务而言有些层并不是直接能用,需要我们微微改一下,比如,resnet 最后的全连接层是分 1000 类,而我们只有 21 类;又比如,resnet 第一层卷积接收的通道是 3, 我们可能输入图片的通道是 4,那么可以通过以下方法修改:

resnet.conv1 = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False) resnet.fc = nn.Linear(2048, 21)

简单预训练

模型已经改完了,接下来我们就进行简单预训练吧。 我们先从 torchvision 中调用基本模型,加载预训练模型,然后,重点来了,将其中的层直接替换为我们需要的层即可

resnet = torchvision.models.resnet152(pretrained=True) # 原本为 1000 类,改为 10 类 resnet.fc = torch.nn.Linear(2048, 10)

其中使用了 pretrained 参数,会直接加载预训练模型,内部实现和前文提到的加载预训练的方法一样。因为是先加载的预训练参数,相当于模型中已经有参数了,所以替换掉最后一层即可。OK!

原文发布于微信公众号 - AI研习社(okweiwu)

原文发表时间:2017-05-02

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏ATYUN订阅号

不用P图软件,TensorFlow pix2pix也能帮你打造出六块腹肌!

可视化结果是十分强大的。然而,在健身领域,要想清楚地看到未来的锻炼结果往往是很困难的。我们是否可以利用深度学习让人们更接近他们的个人健康目标,从而帮助他们设想未...

63440
来自专栏ATYUN订阅号

Keras和PyTorch的视觉识别与迁移学习对比

在上一篇文章中,我们简述了Keras和PyTorch的区别,旨在帮助你选择更适合你需求的框架。现在,我们进行实战进行。我们将让Keras和PyTorch互相较量...

1.2K40
来自专栏专知

【前沿】TensorFlow Pytorch Keras代码实现深度学习大神Hinton NIPS2017 Capsule论文

【导读】10月26日,深度学习元老Hinton的NIPS2017 Capsule论文《Dynamic Routing Between Capsules》终于在a...

43460
来自专栏奇点大数据

Pytorch神器(9)

大家好,今天我们来看Pytorch神器第九次分享的内容——Image Caption。

26220
来自专栏程序员的知识天地

前端工程师深度学习,就能在浏览器里玩转深度学习

TensorFlow.js 的发布可以说是 JS 社区开发者的福音!但是在浏览器中训练一些模型还是会存在一些问题与不同,如何可以让训练效果更好?本文为大家总结了...

18630
来自专栏量化投资与机器学习

使用CNN(LSTM架构)进行序列预测基于TensorFlow

作者 / mouradmourafiq 翻译 / 编辑部翻译组 来源 / https://github.com/mouradmourafiq 前言 这篇推文抛...

65870
来自专栏AI研习社

你在数据预处理上花费的时间,是否比机器学习还要多?

Nuts-ml 是一个新的 Python 数据预处理库,专门针对视觉领域的 GPU 深度学习应用。 它以独立、可复用的单元模块的形式,提供主流数据预处理函数。...

40380
来自专栏专知

【专知国庆特刊-PyTorch手把手深度学习教程系列01】一文带你入门优雅的PyTorch

【导读】主题链路知识是我们专知的核心功能之一,为用户提供AI领域系统性的知识学习服务,一站式学习人工智能的知识,包含人工智能( 机器学习、自然语言处理、计算机视...

1.2K80
来自专栏Python小屋

使用决策树分类算法判断一下你的Python水平

决策树算法是一种常用的机器学习算法,既可以用于分类,也可以用于回归。决策树算法类似于程序设计语言中嵌套的多分支选择结构,根据不同的条件选择不同的分支路径,最终到...

13620
来自专栏尾尾部落

使用自己的语料训练word2vec模型

先对新闻文本进行分词,使用的是结巴分词工具,将分词后的文本保存在seg201708.txt,以备后期使用。

1.1K30

扫码关注云+社区

领取腾讯云代金券