pytorch学习笔记------迁移学习②

迁移学习两种训练策略：

        ①直接用原参数做初始化，继续进行训练

        ②将前面卷积层（主要用于特征提取）冻住，直接训练全连接层，可根据数据量多少调整冻住层数量

pytorch迁移学习步骤代码示例

加载resnet152

import torch
from torchvision import models
import torch.nn as nn
model_name='resnet'
feature_extract=True
device=torch.device('cuda:0'if torch.cuda.is_available() else 'cpu')
model_ft=models.resnet152()
print(model_ft)

 我们可以用以下代码决定是否冻住一些层

'''
迁移学习策略是否冻住哪个层
'''
def set_parameter_requires_grad(model,feature_edxtracting):
    if feature_edxtracting:
        for param in models.parameters():
            param.require_grad=False

为方便加载各种网络我们可以定义一个初始化加载各种经典网络的函数

'''
模型选择init函数
'''
def initialize_model(model_name,num_classes,feature_extract,use_pretrained=True):
    model_ft=None
    input_size=0
    if model_name=='resnet':
        model_ft=models.resnet152(pretrained=use_pretrained)#是否用训练好的模型继续训练
        set_parameter_requires_grad(model_ft,feature_extract)#决定是否冻住一些层
        num_ftrs=model_ft.fc.in_features#原网络fc全连接层
        model_ft.fc=nn.Sequential(nn.Linear(num_ftrs,num_classes),nn.LogSoftmax(dim=1))#更改全连接层输出参数
        input_size=224
    elif model_name=='vgg':
        pass
    else:
        pass
    return model_ft,input_size

如果pretrain为True则表示加载预训练模型会自动下载至此路径

实例化加载网络

import torch
from torchvision import models
import torch.nn as nn
model_name='resnet'
feature_extract=True
device=torch.device('cuda:0'if torch.cuda.is_available() else 'cpu')
model_ft,input_size=initialize_model(model_name,102,feature_extract,use_pretrained=True)
model_ft=model_ft.to(device)
filename='checkpoint.pth'

params_to_update=model_ft.parameters()
if feature_extract:
    params_to_updata=[]
    for name,param in model_ft.named_parameters():
        if param.requires_grad==True:
            params_to_updata.append(param)
            print("\t",name)
else:
    for name,param in model_ft.named_parameters():
        if param.requires_grad==True:
            print("\t",name)