python PyTorch预训练示例

前言

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

直接加载预训练模型

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

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

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

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

还有第二种加载方法：

1	`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，那么可以通过以下方法修改：

1 2	`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！

以上就是本文的全部内容，希望对大家的学习有所帮助，也希望大家多多支持服务器之家。

原文链接：http://blog.csdn.net/AMDS123/article/details/70144935

秒客网

python PyTorch预训练示例

相关文章