利用ImageNet下的预训练权重采用迁移学习策略,能够实现模型快速训练,提高图像分类性能。下面以vgg和resnet网络模型为例,微调最后的分类层进行分类。
(1)微调VGG模型进行图像分类(以vgg16为例)
import torch import torchvision.models as models classes_num = 200 # 数据集的类别数 model = models.vgg16(pretrained=True) model.classifier[-1].out_features = classes_num model = model.cuda() print(model)
(2)微调ResNet模型进行图像分类(以ResNet-34为例)
import torch import torchvision.models as models classes_num = 200 #数据集的类别数 model = models.resnet34(pretrained=True) fc_features = model.fc.in_features model.fc = torch.nn.Linear(fc_features, classes_num) model = model.cuda() print(model)
原文:https://www.cnblogs.com/rs-xiaosheng/p/12892362.html