使用迁移模型并修改部分参数

使用经典的AlexNet模型,导入torchvision库中的models,并利用models。alexnet()函数加载预设的模型,其中参数pretrained=True代表加载经过了训练后的模型参数。

AlexNet分成了features和classifier两大块。其中features模块负责提取特征,以卷积层为主,classifier模块负责分类,以全连接层为主。

为了构造一个二元分类器,需要重新定义AlexNet的classifier模块。前两个全连接层的参数可以保持不变,最后一层输入改成2:

import torch.nn as nn

for param in  alexnet.parameters():
	param.requires_grad = False
alexnet.classifier=nn.Sequential(
	nn.Dropout(),
	nn.Linear(256*6*6,4096),
	nn.ReLU(inplace = True),
	nn.Dropout(),
	nn.Linear(4096,4096),
	nn.ReLU(inplace=True),
	nn.Linear(4096,2),)

循环遍历AlexNet中的所有参数,并将参数的requires_grad设置为False,这样做可以限制这些参数的更新,而重新定义的classifier模块的参数则默认保持requires_grad为True的设置,这就可以保证在之后的迁移学习的过程中,只更新全连接层的参数,而不更新特征提取层的参数。

摘自:《Pytorch深度学习入门》曾芃壹

你可能感兴趣的:(pytorch,深度学习,迁移学习)