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

Posted AI开发者

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了PyTorch 的预训练,是时候学习一下了相关的知识,希望对你有一定的参考价值。


  前言

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

  直接加载预训练模型

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

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研习社按:本文作者 ,文章原载于作者的

延伸阅读:

关注 AI 研习社后,回复【1】获取

【千 G 神经网络 / 人工智能 / 大数据、教程、论文!】

---------------------------------------------

   “TensorFlow & 神经网络算法高级应用班” 要开课啦!

从初级到高级,理论 + 实战,一站式深度了解 TensorFlow!

本课程面向深度学习开发者,讲授如何利用 TensorFlow 解决图像识别、文本分析等具体问题。课程跨度为 10 周,将从 TensorFlow 的原理与基础实战技巧开始,一步步教授学员如何在 TensorFlow 上搭建 CNN、自编码、RNN、GAN 等模型,并最终掌握一整套基于 TensorFlow 做深度学习开发的专业技能。

两名授课老师佟达、白发川身为 ThoughtWorks 的资深技术专家,具有丰富的大数据平台搭建、深度学习系统开发项目经验。

时间:每周二、四晚 20:00-21:00

开课时长:总学时 20 小时,分 10 周完成,每周 2 次,每次 1 小时

以上是关于PyTorch 的预训练,是时候学习一下了的主要内容,如果未能解决你的问题,请参考以下文章

批量大小降低了预训练 CNN 集合的准确性

PyTorch下载的预训练模型的保存位置(Windows)

论文泛读171具有对抗性扰动的自监督对比学习,用于鲁棒的预训练语言模型

论文泛读171具有对抗性扰动的自监督对比学习,用于鲁棒的预训练语言模型

Pytorch - 跳过计算每个时期的预训练模型的特征

有啥方法可以将 PyTorch 中可用的预训练模型下载到特定路径?