使用经典的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深度学习入门》曾芃壹
|