Pytorch下微调网络模型（迁移学习）进行图像分类

Pytorch下微调网络模型进行图像分类
利用ImageNet下的预训练权重采用迁移学习策略，能够实现模型快速训练，提高图像分类性能。下面以vgg和resnet网络模型为例，微调最后的分类层进行分类。

注意，微调只对分类层(也就是全连接层)的参数进行更新，前面的参数需要被冻结。

（1）微调VGG模型进行图像分类（以vgg16为例）

import torch
import torch.nn as nn
import torchvision.models as models

classes_num = 200 # 数据集的类别数

model = models.vgg16(pretrained=True) # vgg16_bn, 区别，有无BN操作
for parameter in model.parameters():
    parameter.required_grad = False
model.classifier = nn.Sequential(nn.Linear(512*7*7, 4096),
                                 nn.ReLU(inplace=True),
                                 nn.Dropout(0.5),
                                 nn.Linear(4096, 4096),
                                 nn.ReLU(inplace=True),
                                 nn.Dropout(0.5),
                                 nn.Linear(4096, classes_num))
model = model.cuda()
print(model)

(2)微调ResNet模型进行图像分类(以ResNet-34为例)

import torch
import torch.nn as nn
import torchvision.models as models

classes_num = 200 # 数据集的类别数

model = models.resnet34(pretrained=True)
for parameter in model.parameters():
    parameter.required_grad = False
model.classifier = nn.Linear(512, classes_num)
model = model.cuda()
print(model)

参考：
https://www.cnblogs.com/rs-xiaosheng/p/12892362.html

实例：
https://blog.csdn.net/heiheiya/article/details/103028543
https://www.jianshu.com/p/b935e108ba7d
https://blog.csdn.net/m0_37867091/article/details/107450768