[英]PyTorch transfer learning with pre-trained ImageNet model
我想在已经在ImageNet上训练的模型上使用转移学习来创建图像分类器。
如何用我自己的自定义分类器替换torchvision.models
ImageNet分类器的最后一层?
获得预先训练的ImageNet模型( resnet152
具有最佳精度):
from torchvision import models
# https://pytorch.org/docs/stable/torchvision/models.html
model = models.resnet152(pretrained=True)
打印出它的结构,以便我们可以比较最终状态:
print(model)
从模型中删除最后一个模块(通常是一个完全连接的层):
classifier_name, old_classifier = model._modules.popitem()
冻结模型的特征检测器部分的参数,以便它们不通过反向传播进行调整:
for param in model.parameters():
param.requires_grad = False
创建一个新的分类器:
classifier_input_size = old_classifier.in_features
classifier = nn.Sequential(OrderedDict([
('fc1', nn.Linear(classifier_input_size, hidden_layer_size)),
('activation', nn.SELU()),
('dropout', nn.Dropout(p=0.5)),
('fc2', nn.Linear(hidden_layer_size, output_layer_size)),
('output', nn.LogSoftmax(dim=1))
]))
我们的分类器的模块名称必须与删除的模块名称相同。 将我们的新分类器添加到功能检测器的末尾:
model.add_module(classifier_name, classifier)
最后,打印出新网络的结构:
print(model)
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.