Pytorch Transfer Learning

PyTorch
所谓迁移学习，就是能让现有的模型算法稍加调整即可应用于一个新的领域和功能的一项技术。

1、加载训练好的模型

import torchvision
# Gives easier dataset managment and creates mini batches
from torch.utils.data import DataLoader
# Has standard datasets we can import in a nice way
import torchvision.datasets as datasets 
# Transformations we can perform on our dataset
import torchvision.transforms as transforms  
# Load pretrain model & modify it
model = torchvision.models.vgg16(pretrained=True)
print(model)

2、特征层处理

train_dataset = datasets.CIFAR10(
    root="dataset/", train=True, transform=transforms.ToTensor(), download=True
)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
model2 = model.features
model2.eval()
train = next(iter(train_loader))
print(len(train))
x = train[0].to(device)
y = train[1].to(device)
y_pred = model2(x) 
print(y_pred.shape)
# torch.Size([64, 512, 1, 1])

结果是在所有的conv层之后，avgpool层的输入是1x1(即512x1x1) (avgpool工作的维度，即1x1)，所以这里不需要执行avgpool(理论上可以应用大小为1x1的avgpool，而不是使用自定义Identity层，结果是一样的)。
顺便说一下，avgpool的参数是“output_size”，而不是“kernel_size”。这意味着它保证最终输出是请求的大小。在这种情况下，1x1将被复制为7x7，这没有任何好处。

class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()
    def forward(self, x):
        return x

因此可以训练以下模型：

# Load pretrain model & modify it
model = torchvision.models.vgg16(pretrained=True)
model.avgpool = Identity()
model.classifier = nn.Linear(512,10)
# 如果只想改变classifier中的某一层，
# 可用model.classifier[i]=nn.Linear() 或者nn.Relu()，
# i是某层所在的索引，本模型中classifier有7层

结果如下：

3、训练部分模型

# 如果不想训练整个模型，可将模型的参数的梯度设置为False。
for param in model.parameters():
    param.requires_grad = False
# 同时对需要训练的部分重新设置网络。
model.avgpool = Identity()
model.classifier = nn.Sequential(
    nn.Linear(512, 100), nn.ReLU(), nn.Linear(100, num_classes)
)