迁移学习模型

发表于 2022-04-19 更新于 2022-08-27 分类于计算机视觉阅读次数： Valine：
本文字数： 9.5k 阅读时长 ≈ 9 分钟
import os
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
import torchvision
import torch.optim as optim
from torchvision import transforms,models,datasets
import imageio
import time
import warnings
import random
import sys
import copy
import json
from PIL import Image

# 数据读取与预处理操作
data_dir = './flower_data/'
train_dir = data_dir + '/train'
valid_dir = data_dir + '/valid'
num_classes = 102

# 制作好数据源

# 数据增强(Data Augmentation)
# 数据不够怎么办？如何更高效的利用数据？
# 对图像进行镜像，翻转，旋转，放大，缩小

# 制作好数据源
# data_transforms中指定了所有图像预处理操作
# ImageFolder假设所有的文件按文件夹保存好，每个文件夹下面存同一类别的图片，文件夹的名字为分类的名字
data_transfroms = {
    'train':transforms.Compose([ # transforms.Compose()类用来组合多个torchvision.transforms操作。
        transforms.RandomRotation(45),# 随机旋转，-45到45度之间随机选
        transforms.CenterCrop(224), #从中心开始裁剪  留下224*224的图像区域   另外还有随机裁剪
        transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转，选择一个概率  做镜像
        transforms.RandomVerticalFlip(p=0.5), # 随机垂直翻转   做镜像
        transforms.ColorJitter(brightness=0.2,contrast=0.1,saturation=0.1,hue=0.1), # 参数1为亮度，参数2为对比度，参数3为饱和度，参数4为色相
        transforms.RandomGrayscale(p=0.025), # 概率转换成灰度率，3通道就是R=G=B
        transforms.ToTensor(), # ???归一化
        # PIL Image或numpy.ndarray(形状为HxWxC)数据范围是[0,255]到一个Torch.FloatTensor

        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]) # 均值，标准差
        # 为了使迁移学习效果更好，用人家模型的均值和标准差
    ]),
    # 除了数据增强外，训练集和验证集的数据处理要一样
    'valid':transforms.Compose([
        transforms.Resize(256), #调整图片的大小，使图片
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]) # 均值，标准差 和训练时参数设置一样
        # 因为这里是预测图像的结果，不需要考虑训练时的一些参数设置
    ]),
}

# Batch数据制作
batch_size = 8 #根据内存适当的调整大小

# datasets.ImageFolder(root,transform,...)
# root 图片存储的根路径 transform 图片进行的变换
# 返回值 self.classes：用一个list保存类别名称;...
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir,x),data_transfroms[x]) for x in ['train','valid']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x],batch_size=batch_size,shuffle=True) for x in ['train','valid']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train','valid']}
class_names = image_datasets['train'].classes
print(class_names)
# print(dataset_sizes)

# 读取标签对应的名字
#  with as 操作已经打开的文件对象（本身就是上下文管理器），无论期间是否抛出异常，都能保证 with as 语句执行完毕后自动关闭已经打开的文件。
# target 参数用于指定一个变量，该语句会将 expression 指定的结果保存到该变量中
with open('cat_to_name.json','r') as f:
    cat_to_name = json.load(f)  # json.load读取文件类型，返回的数据类型是字典
    print(cat_to_name)

# 展示下数据
# 注意tensor的数据需要转化成numpy的格式，而且还需要还原回标准化的结果
def im_convert(tensor):
    # 展示数据
    image = tensor.to("cpu").clone().detach()
    image = image.numpy().squeeze()
    image = image.transpose(1,2,0) # h,w,c     在tensor中c,h,w
    image = image * np.array((0.229,0.224,0.225)) + np.array((0.485,0.456,0.406)) # 把数据还原回去
    image = image.clip(0,1)
    return image

fig = plt.figure(figsize=(20,12))
# plt.figure设置画布的大小和背景   figsize:设置画布大小
columns = 4
rows = 2

dataiter = iter(dataloaders['train']) # iter的输入是支持迭代的集合对象
print(dataiter)
inputs,classes = dataiter.next()
print(classes)

for idx in range(columns * rows):
    # add_subplot(rows, columns, idx) 将画布分成rows行，columns列，图像画在从左到右从上到下的第idx块
    ax = fig.add_subplot(rows, columns, idx+1, xticks=[], yticks=[])
    ax.set_title(cat_to_name[str(int(class_names[classes[idx]]))])
    plt.imshow(im_convert(inputs[idx]))
# plt.show()

# 加载models中提供的模型，并且直接用训练好的权重当做初试化参数
# 第一次执行需要下载，可能会比较慢

model_name = 'resnet' # 可选的比较多['ResNet','alexnet','vgg','squeezenet','densenet','inception']
# 是否用人家训练好的特征来做
feature_extract = True
# 是否用GPU训练
train_on_gpu = torch.cuda.is_available()
if not train_on_gpu:
    print('GPU false')
else:
    print('GPU True')

#torch.device包含一个设备类型（'cpu'或'cuda'设备类型）和可选的设备的序号
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# print(device)

# 设置不更新的部分
def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

model_ft = models.resnet152()
# print(model_ft)

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) # pretrained为true 会下载训练好的模型
        set_parameter_requires_grad(model_ft,feature_extract) # 选择性的冻住哪些层
        num_ftrs = model_ft.fc.in_features # 用fc.in_features得到最后一层的输出，即下一层的输入
        # 重写fc层
        # softmax和logsoftmax
        # softmax是 该数求指数 / 所有数求指数的和
        # logsoftmax是 log(softmax) 把softmax的结果log一下
        model_ft.fc = nn.Sequential(nn.Linear(num_ftrs,num_classes),nn.LogSoftmax(dim=1))
        input_size = 224

    return model_ft, input_size

# 设置哪些层需要训练
model_ft, input_size = initialize_model(model_name,num_classes,feature_extract,use_pretrained=True)
# GPU计算
# 构造好的模型加入GPU
model_ft = model_ft.to(device)

# 模型保存
filename = 'checkpoint.pth'

# 是否训练所有层
params_to_update = model_ft.parameters()
print('Params to learn:')
# 这个if else是为了查看哪些参数是需要被训练的
if feature_extract:
    params_to_update = []
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            params_to_update.append(param)
            print("\t",name)
else:
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            print("\t",name)

# 优化器设置
optimizer_ft = optim.Adam(model_ft.parameters(),lr=1e-2)
scheduler = optim.lr_scheduler.StepLR(optimizer_ft,step_size=7,gamma=0.1) # 学习率每7个epoch衰减成原来的1/10
# 最后一层已经LogSoftmax()了，所以不能nn.CrossEntropyLoss()来计算了，nn.CrossEntropyLoss()相当于logSoftmax()和nn.NLLLoss()整合
criterion = nn.NLLLoss()

# 训练模型
def train_model(model, dataloaders, criterion, optimizer, num_epochs=10, is_inception=False, filename=filename):
    since = time.time()
    best_acc = 0 # 最好的准确率
    # ！！！！不是保存最后一个模型，而是保存准确率较高的模型

    # 训练前把模型加入GPU
    model.to(device)

    val_acc_history = []
    train_acc_history = []
    train_losses = []
    valid_losses = []

    # optimizer.param_groups： 是长度为2的list，其中的元素是2个字典；
    # optimizer.param_groups[0]： 长度为6的字典，包括[‘amsgrad’, ‘params’, ‘lr’, ‘betas’, ‘weight_decay’, ‘eps’]这6个参数；
    # optimizer.param_groups[1]： 好像是表示优化器的状态的一个字典；
    LRs = [optimizer.param_groups[0]['lr']] #学习率数组
    # copy.deepcopy() python中的函数进行深拷贝
    best_model_wts = copy.deepcopy(model.state_dict())

    for epoch in range(num_epochs):
        print("Epoch {}/{}".format(epoch,num_epochs-1))
        print('-'*10)

        # 训练和验证
        for phase in ['train','valid']:
            if phase == 'train':
                model.train() # 训练
            else:
                model.eval() # 验证

            running_loss = 0.0
            running_corrects = 0

            # 把数据都取个遍
            for inputs,labels in dataloaders[phase]:
                # 把训练的数据加入GPU
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 清零
                optimizer.zero_grad()
                # 只有训练的时候计算和更新梯度
                    # set_grad_enabled  用于设置梯度计算打开或关闭状态的上下文管理器.
                with torch.set_grad_enabled(phase == 'train'):
                    if is_inception and phase == 'train':
                        outputs, aux_outputs = model(inputs)
                        loss1 = criterion(outputs,labels)
                        loss2 = criterion(aux_outputs,labels)
                        loss = loss1 + 0.4*loss2
                    else:
                        outputs = model(inputs)
                        loss = criterion(outputs,labels)
                        # print(type(outputs))
                    _,preds = torch.max(outputs,1)

                    # 训练阶段更新权重
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # 计算损失
                running_loss +=loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            time_elapsed = time.time() - since
            print('Time elapsed {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

            # 得到最好那次的模型
            if phase == 'valid' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                # pytorch中model.parameters()和model.state_dict()使用时的区别
                # model.parameters()方法返回的是一个生成器generator，里面只有模型的参数，而没有对应的网络层名称；
                # 并且 如果有些层的参数没有被训练，那么它只保存那些被训练的参数
                # model.state_dict() 返回的则是一个字典 {key:value}，key 是网络层名称，value 则是该层的参数

                # state_dict()是一个状态字典,⼀个简单的python的字典对象,将每⼀层与它的对应参数建⽴映射关系
                state = {
                    'state_dict': model.state_dict(),
                    'best_acc': best_acc,
                    # 优化器参数：有时保存模型的参数需要稍后接着训练，那么就必须保存优化器的状态和其所使用的超参数
                    'optimizer': optimizer.state_dict(), # 包含了优化器的状态，以及被使用的超参数
                }
                torch.save(state, filename)
            if phase == 'valid':
                val_acc_history.append(epoch_acc)
                valid_losses.append(epoch_loss)
                # scheduler.step（）按照Pytorch的定义是用来更新优化器的学习率的，一般是按照epoch为单位进行更换，即多少个epoch后更换一次学习率
                # ???epoch_loss这个参数有什么意义呢
                scheduler.step(epoch_loss)
            if phase == 'train':
                train_acc_history.append(epoch_acc)
                train_losses.append(epoch_loss)

        print('Optimizer learning rate : {:.7f}'.format(optimizer.param_groups[0]['lr']))
        LRs.append(optimizer.param_groups[0]['lr'])
        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # 训练完后用最好的一次当做模型最终的结果
    model.load_state_dict(best_model_wts)
    return model, val_acc_history, train_acc_history, valid_losses, train_losses, LRs

model_ft, val_acc_history, train_acc_history, valid_losses, train_losses, LRs  = train_model(model_ft, dataloaders, criterion, optimizer_ft, num_epochs=10, is_inception=(model_name=="inception"))