pytorch vgg16& resnet50简单的迁移学习

 

import torch
import torch.nn as nn
from torch.utils.data import dataset,dataloader,Dataset,DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torchvision.models import vgg16

vgg = vgg16(pretrained=True)
#固定所有参数，requires_grad=False
for param in vgg.parameters():
    param.requires_grad = False

vgg.classifier[6] = nn.Sequential(nn.Linear(in_features=4096,out_features=22))

for p in vgg.parameters():
    print(p)

或者直接使用以下代码只更新最后一层参数：

for param in model.classifier[6].parameters():
    param.requires_grad = True

 

#resnet

resnet = models.resnet50(pretrained=True)
for param in resnet.parameters():      # 冻结参数
    param.requires_grad = False
resnet.fc = nn.Linear(resnet.fc.in_features, 100)