Pytorch实现CNN模型的迁移学习——蜜蜂和蚂蚁图片分类项目

很多时候当训练一个新的图像分类任务时，一般不会完全从一个随机的模型开始训练，而是利用预训练的模型来加速训练的过程。经常使用在ImageNet上的预训练模型。

• 这是一种transfer learning的方法。我们常用以下两种方法做迁移学习。
  • fine tuning: 从一个预训练模型开始，我们改变一些模型的架构，然后继续训练整个模型的参数。
  • feature extraction: 我们不再改变与训练模型的参数，而是只更新我们改变过的部分模型参数。我们之所以叫它feature extraction是因为我们把预训练的CNN模型当做一个特征提取模型，利用提取出来的特征做来完成我们的训练任务。
    以下是构建和训练迁移学习模型的基本步骤：
• 初始化预训练模型
• 把最后一层的输出层改变成我们想要分的类别总数
• 定义一个optimizer来更新参数
• 模型训练

一、项目介绍：

1、目标：本文尝试采用CNN实现图像蜜蜂和蚂蚁图像分类任务。
2、数据说明：使用hymenoptera_data数据集。包括两类图片, bees 和 ants, 这些数据都被处理成了可以使用ImageFolder来读取的格式。
输入数据维度：；输出数据维度
3、torchvision的datasets.ImageFolder参数说明：
（1）data_dir：数据的存储目录。设置成数据的根目录
（2）model_name：训练时使用的模型。设置成自己的训练模型，也可以使用封装好的模型。如，resnet, alexnet, vgg, squeezenet, densenet, inception等
（3）num_classes：表示数据集分类的类别数
（4）batch_size
（5）num_epochs
（6）feature_extract：表示训练时使用fine tuning还是feature extraction方法。如果feature_extract = False，整个模型都会被同时更新。如果feature_extract = True，只有模型的最后一层被更新。
4、网络框架：

二、设置参数

import numpy as np
import torchvision
from torchvision import datasets, transforms, models

import matplotlib.pyplot as plt
import time
import os
import copy
print("Torchvision Version: ",torchvision.__version__)

data_dir = './hymenoptera_data'
#模型可以选择[resnet, alexnet, vgg, squeezenet, densenet, inception]
model_name = 'resnet'
#数据的label分类个数
num_classes = 2
#训练的batch_size
batch_size=32
#训练的轮数
num_epochs = 15
#整个模型参数都参数训练，进行更新
feature_extract = True

二、代码实现

1、初始化模型

#定义需要更新的参数
def set_parameter_requires_grad(model,feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False
#使用resnet框架进行模型初始化           
def initialize_model(model_name,num_class,feature_extract,use_pretrained=True):
    if model_name == 'renet':
        model_ft = models.resnet18(pretrained=use_pretrained)  #使用resnet18最为初始化模型
        set_parameter_requires_grad(model_ft,feature_extract)  #模型中所有参数都更新
        num_ftrs = model_ft.fc.in_features
        model_ft.fc = nn.Linear(num_ftrs,num_classes)  #根据分类个数重新定义最后一层全连接层
        input_size = 224
        
    return model_ft,input_size
model_ft,input_size = initialize_model(model_name,num_classes,feature_extract,use_pretrained=True)
print(model_ft) 

2、导入数据集

根据模型输入的size,将数据预处理称为对应的格式

#定义图片的处理方式
data_transforms = {
    'train':transforms.Compose([
        transforms.RandomResizedCrop(input_size),
        transforms.RandomHorizontalFlip(),
        tranforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  #对三个通道分别从用不同的形式进行标准化
    ]),
    'val':transforms.Compose([
        transforms.Resize(input_size),
        transforms.CenterCrop(input_size),
        transforms.ToTensor(),
        transforms.Normalize([[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]])
    ])
}

print("Initializing Datasets and Dataloaders...")

#生成训练集和测试集
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir,x),data_transforms[x]) for x in ['train','val']}
#生成训练集和测试集的迭代器
dataloaders_dict = {x: torch.utils.data.DataLoader(image_datasets[x],batch_size=batch_size,shuffle=True,num_workers=4) for x in ['train','val']}

#判断是否使用gpu
device = torch.device('cuda:0' if torch.cuda.is_availabel() else 'cpu')

3、观察原始数据

it = iter(dataloaders_dict["train"])
inputs, labels = next(it)
for inputs, labels in dataloaders_dict["train"]:
    print(labels.size())
    
print(len(dataloaders_dict["train"].dataset.imgs),len(dataloaders_dict["train"].dataset))

4、模型训练

torch.manual_seed(53113
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
                  
def train_model(model,dataloaders,criterion,optimizer,num_epochs=5):
    since = time.time()
    val_acc_history = []
    best_model_wts = copy.deepcopy(model.state_dict())   #保存最佳参数；copy.deepcopy一旦复制出来了，就应该是独立的了，原始变量改变其不会变，数值一样，但对象不同的。
    best_acc = 0.  #保存最高的准确率
    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs-1}')
        print("-"*10)
        
        for phase in ['train','val']:
            running_loss = 0.
            running_corrects = 0.
            if phase == 'train':
                model.train()
            else:
                model.eval()
                
            for inputs, labels in dataloaders[phase]:   
                inputs = inputs.to(device)
                labels = labels.to(device)
                
                with torch.autograd.set_grad_enabled(phase=='train'):   #当phase=='train'为True是，可以进行自动求梯度计算。
                    outputs = model(inputs)   #将数据输入到模型中进行预测
                    loss = criterion(outputs,labels)  #损失计算
                    
                _,preds = torch.max(outputs,1)   #在所有输出中选取最大的最为输出概率
                if phase == 'train':
                    optimizer.zero_grad()   #每一次梯度计算后，所有梯度设置为0
                    loss.backward()    #反向求导
                    optimizer.step()   #参数优化
                    
                running_loss += loss.item() * inputs.size(0)   #损失*batch，等于每一批的总损失；每一批的总损失相加为所有样本的损失和
                running_corrects += torch.sum(preds.view(-1) == labels.view(-1))  #每批的预测正确的样本数量；每批预测正确的样本数量相加为所有样本中预测正确的样本数量
            
            epoch_loss = running_loss / len(dataloaders[phase].dataset)  #样本总损失 / 样本总个数 = 平均损失
            epoch_acc = running_corrects / len(dataloaders[phase].dataset) #样本预测正确的个数 / 样本总个数 = 平均准确率
                
            print(f'{phase} Loss:{epoch_loss} Acc:{epoch_acc}')  #取平均准确率最为当前epoch的准确率，与之前的进行比较，保存最高的
            if phase == 'val' and epoch_acc > best_acc:
                  best_acc = epoch_acc  #每个epoch更新一次准确率
                  best_model_wts = copy.deepcopy(model.state_dict())
            if phase == 'val':
                  val_acc_history.append(epoch_acc)  #将每个epoch的准确率添加到列表中
 
            
    time_elapsed = time.time() - since
    print(f'Training compete in {time_elapsed // 60}m {time_elapsed % 60}s')   #time_elapsed // 60表示相除后取商。
    print(f'Best val Acc: {best_acc}')
    
    model.load_state_dict(best_model_wts)
    return model,val_acc_history   #返回模型和评估集，每个epoch的最佳acc    

（1）对参数进行训练的训练集、测试集结果

model_ft = model_ft.to(device)

params_to_update = model_ft.parameters()  #获取模型所有的参数进行更新
print('Params to learn:')
if feature_extract:  #如果feature_extract为True,表示对参数进行更新
    params_to_update = []
    for name,param in model_ft.named_parameters():
        if param.requirs_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)

训练结果：
Training compete in 0.0m 14.700076580047607s
Best val Acc: 0.954248366013072
（2）不对参数训练的训练集和测试集结果

scratch_model,_ = initlialize_model(model_name,num_classes,feature_extract=False,use_pretrained=False)
scratch_model = scratch_model.to(device)
scratch_optimizer = optim.SGD(scratch_model.parameters(), lr=0.001, momentum=0.9)
scratch_criterion = nn.CrossEntropyLoss()
-,scratch_hist = train_model(scratch_model, dataloaders_dict, scratch_criterion, scratch_optimizer, num_epochs=num_epochs)

训练结果：
Training compete in 0.0m 18.418847799301147s
Best val Acc: 0.7450980392156863
（3）是否对参数训练，效果对比

plt.title("Validation Accuracy vs. Number of Training Epochs")
plt.xlabel("Training Epochs")
plt.ylabel("Validation Accuracy")
plt.plot(range(1,num_epochs+1),ohist,label="Pretrained")
plt.plot(range(1,num_epochs+1),scratch_hist,label="Scratch")
plt.ylim((0,1.))
plt.xticks(np.arange(1, num_epochs+1, 1.0))
plt.legend()
plt.show()

参考文献：七月在线Pytorch课程