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)