PyTorch 图像：02.微调 Torchvision 模型

在本教程中，我们将深入探讨如何对 torchvision 模型进行微调和特征提取，所有这些模型都已经预先在1000类的magenet数据集上训练完成。 本教程将深入介绍如何使用几个现代的CNN架构，并将直观展示如何微调任意的PyTorch模型。由于每个模型架构是有差异的，因此没有 可以在所有场景中使用的微调代码样板。然而，研究人员必须查看现有架构并对每个模型进行自定义调整。

在本文档中，我们将执行两种类型的转移学习：微调和特征提取。

在微调中，我们从预训练模型开始，更新我们新任务的所有模型参数，实质上是重新训练整个模型。

在特征提取中，我们从预训练模型开始，仅更新从中导出预测的最终图层权重。它被称为特征提取，因为我们使用预训练的CNN作为固定 的特征提取器，并且仅改变输出层。

有关迁移学习的更多技术信息，请参阅此处和这里。

通常，这两种迁移学习方法都遵循以下几个步骤：

初始化预训练模型

重组最后一层，使其具有与新数据集类别数相同的输出数

为优化算法定义我们想要在训练期间更新的参数

运行训练步骤

1.导入相关包并打印版本号

from __future__ import print_function
from __future__ import division
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
print("PyTorch Version: ",torch.__version__)
print("Torchvision Version: ",torchvision.__version__)

输出结果：

PyTorch Version:  1.1.0
Torchvision Version:  0.3.0

2.输入

以下为运行时需要更改的所有参数。我们将使用的数据集 hymenoptera_data 可在此处 下载。该数据集包含两类：蜜蜂和蚂蚁，其结构使得我们可以使用 ImageFolder 数据集，不需要编写我们自己的自定义数据集。 下载数据并设置data_dir为数据集的根目录。model_name是您要使用的模型名称，必须从此列表中选择：

[resnet, alexnet, vgg, squeezenet, densenet, inception]

其他输入如下：num_classes为数据集的类别数，batch_size是训练的 batch 大小，可以根据您机器的计算能力进行调整，num_epochsis是 我们想要运行的训练 epoch 数，feature_extractis是定义我们选择微调还是特征提取的布尔值。如果feature_extract = False， 将微调模型，并更新所有模型参数。如果feature_extract = True，则仅更新最后一层的参数，其他参数保持不变。

# 顶级数据目录。 这里我们假设目录的格式符合ImageFolder结构
data_dir = "./data/hymenoptera_data"
# 从[resnet, alexnet, vgg, squeezenet, densenet, inception]中选择模型
model_name = "squeezenet"
# 数据集中类别数量
num_classes = 2
# 训练的批量大小（根据您的内存量而变化）
batch_size = 8
# 你要训练的epoch数
num_epochs = 15
# 用于特征提取的标志。 当为False时，我们微调整个模型，
# 当True时我们只更新重新形成的图层参数
feature_extract = True

3.辅助函数

在编写调整模型的代码之前，我们先定义一些辅助函数。

3.1 模型训练和验证代码

train_model函数处理给定模型的训练和验证。作为输入，它需要PyTorch模型、数据加载器字典、损失函数、优化器、用于训练和验 证epoch数，以及当模型是初始模型时的布尔标志。is_inception标志用于容纳 Inception v3 模型，因为该体系结构使用辅助输出， 并且整体模型损失涉及辅助输出和最终输出，如此处所述。 这个函数训练指定数量的epoch,并且在每个epoch之后运行完整的验证步骤。它还跟踪最佳性能的模型（从验证准确率方面），并在训练 结束时返回性能最好的模型。在每个epoch之后，打印训练和验证正确率。

def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, is_inception=False):
   since = time.time()
   val_acc_history = []
   best_model_wts = copy.deepcopy(model.state_dict())
   best_acc = 0.0
   for epoch in range(num_epochs):
       print('Epoch {}/{}'.format(epoch, num_epochs - 1))
       print('-' * 10)
       # 每个epoch都有一个训练和验证阶段
       for phase in ['train', 'val']:
           if phase == 'train':
               model.train()  # Set model to training mode
           else:
               model.eval()   # Set model to evaluate mode
           running_loss = 0.0
           running_corrects = 0
           # 迭代数据
           for inputs, labels in dataloaders[phase]:
               inputs = inputs.to(device)
               labels = labels.to(device)
               # 零参数梯度
               optimizer.zero_grad()
               # 前向
               # 如果只在训练时则跟踪轨迹
               with torch.set_grad_enabled(phase == 'train'):
                   # 获取模型输出并计算损失
                   # 开始的特殊情况，因为在训练中它有一个辅助输出。
                   # 在训练模式下，我们通过将最终输出和辅助输出相加来计算损耗
                   # 但在测试中我们只考虑最终输出。
                   if is_inception and phase == 'train':
                       # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958
                       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)
                   _, preds = torch.max(outputs, 1)
                   # backward + optimize only if in training phase
                   if phase == 'train':
                       loss.backward()
                       optimizer.step()
               # 统计
               running_loss += loss.item() * inputs.size(0)
               running_corrects += torch.sum(preds