(β) 计算机视觉的量化迁移学习教程

创建时间：2019年12月06日 | 最后更新时间：2021年07月27日 | 最后验证时间：2024年11月05日

提示

要充分利用本教程，我们建议使用此 Colab 版本。 这将允许您对下面呈现的信息进行实验。

作者: Zafar Takhirov

审阅人: Raghuraman Krishnamoorthi

编辑者: 林嘉诚

本教程基于由 Sasank Chilamkurthy 编写的原始 PyTorch 迁移学习 教程进行构建。

迁移学习指的是利用预训练模型来应用于不同数据集的技术。 迁移学习主要有两种使用方式：

1. ConvNet 作为固定的特征提取器：在这里，您“冻结” 网络中所有参数的权重，除了最后几层（也称为“头部”，通常是全连接层）。 这些最后的层被替换为用随机权重初始化的新层，并且只训练这些层。

2. 微调卷积网络: 与随机初始化不同，模型使用预训练网络进行初始化，之后训练过程如常进行，但使用不同的数据集。 通常在网络中也会替换头部（或部分头部），以防输出数量不同。 在该方法中，通常会将学习率设置为较小的数值。 这是因为在网络已经训练好的情况下，只需进行少量调整即可“微调”它以适应新的数据集。

你也可以结合上述两种方法： 首先你可以冻结特征提取器，并训练头部。之后，你可以解冻特征提取器（或其中一部分），将学习率设置为更小的值，并继续训练。

在这一部分，你将使用第一种方法 – 使用量化模型提取特征。

第0部分. 先决条件

在深入探讨迁移学习之前，让我们先回顾一下“前提条件”，例如安装和数据加载/可视化。

# Imports
import copy
import matplotlib.pyplot as plt
import numpy as np
import os
import time

plt.ion()

安装夜间构建版本

因为你将使用PyTorch的测试版部分，建议安装最新版本的torch和torchvision。你可以在这里找到最新的本地安装说明。 例如，要安装不带GPU支持的版本：

pip install numpy
pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
# For CUDA support use https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html

加载数据

注意

本节与原始的迁移学习教程完全相同。

我们将使用 torchvision 和 torch.utils.data 包来加载 数据。

你今天要解决的问题是从图像中分类蚂蚁和 蜜蜂。数据集包含约120张用于蚂蚁和蜜蜂的训练图像。 每个类别有75张验证图像。这被认为是一个非常小的数据集，难以进行泛化。 然而，由于我们使用了迁移学习，我们应该能够合理地进行泛化。

这个数据集是imagenet的一个非常小的子集。

注意

从这里下载数据，并将其解压到data目录。

import torch
from torchvision import transforms, datasets

# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
    'train': transforms.Compose([
        transforms.Resize(224),
        transforms.RandomCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = 'data/hymenoptera_data'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=16,
                                              shuffle=True, num_workers=8)
              for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

可视化几幅图像

让我们可视化一些训练图像，以便理解数据增强。

import torchvision

def imshow(inp, title=None, ax=None, figsize=(5, 5)):
  """Imshow for Tensor."""
  inp = inp.numpy().transpose((1, 2, 0))
  mean = np.array([0.485, 0.456, 0.406])
  std = np.array([0.229, 0.224, 0.225])
  inp = std * inp + mean
  inp = np.clip(inp, 0, 1)
  if ax is None:
    fig, ax = plt.subplots(1, figsize=figsize)
  ax.imshow(inp)
  ax.set_xticks([])
  ax.set_yticks([])
  if title is not None:
    ax.set_title(title)

# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))

# Make a grid from batch
out = torchvision.utils.make_grid(inputs, nrow=4)

fig, ax = plt.subplots(1, figsize=(10, 10))
imshow(out, title=[class_names[x] for x in classes], ax=ax)

模型训练支持功能

以下是用于模型训练的通用函数。 该函数也

• 安排学习率

• 保存最佳模型

def train_model(model, criterion, optimizer, scheduler, num_epochs=25, device='cpu'):
  """
  Support function for model training.

  Args:
    model: Model to be trained
    criterion: Optimization criterion (loss)
    optimizer: Optimizer to use for training
    scheduler: Instance of ``torch.optim.lr_scheduler``
    num_epochs: Number of epochs
    device: Device to run the training on. Must be 'cpu' or 'cuda'
  """
  since = time.time()

  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)

    # Each epoch has a training and validation phase
    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

      # Iterate over data.
      for inputs, labels in dataloaders[phase]:
        inputs = inputs.to(device)
        labels = labels.to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward
        # track history if only in train
        with torch.set_grad_enabled(phase == 'train'):
          outputs = model(inputs)
          _, preds = torch.max(outputs, 1)
          loss = criterion(outputs, labels)

          # backward + optimize only if in training phase
          if phase == 'train':
            loss.backward()
            optimizer.step()

        # statistics
        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)
      if phase == 'train':
        scheduler.step()

      epoch_loss = running_loss / dataset_sizes[phase]
      epoch_acc = running_corrects.double() / dataset_sizes[phase]

      print('{} Loss: {:.4f} Acc: {:.4f}'.format(
        phase, epoch_loss, epoch_acc))

      # deep copy the model
      if phase == 'val' and epoch_acc > best_acc:
        best_acc = epoch_acc
        best_model_wts = copy.deepcopy(model.state_dict())

    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))

  # load best model weights
  model.load_state_dict(best_model_wts)
  return model

模型预测可视化支持功能

通用函数以显示几幅图像的预测结果

def visualize_model(model, rows=3, cols=3):
  was_training = model.training
  model.eval()
  current_row = current_col = 0
  fig, ax = plt.subplots(rows, cols, figsize=(cols*2, rows*2))

  with torch.no_grad():
    for idx, (imgs, lbls) in enumerate(dataloaders['val']):
      imgs = imgs.cpu()
      lbls = lbls.cpu()

      outputs = model(imgs)
      _, preds = torch.max(outputs, 1)

      for jdx in range(imgs.size()[0]):
        imshow(imgs.data[jdx], ax=ax[current_row, current_col])
        ax[current_row, current_col].axis('off')
        ax[current_row, current_col].set_title('predicted: {}'.format(class_names[preds[jdx]]))

        current_col += 1
        if current_col >= cols:
          current_row += 1
          current_col = 0
        if current_row >= rows:
          model.train(mode=was_training)
          return
    model.train(mode=was_training)

第1部分. 基于量化特征提取器训练自定义分类器

在本节中，您将使用“冻结”的量化特征提取器，并在其上训练一个自定义的分类器头部。与浮点模型不同，您不需要为量化模型设置 requires_grad=False，因为它没有可训练的参数。请参阅 文档 获取更多详情。

加载一个预训练模型：对于这个练习，您将使用 ResNet-18。

import torchvision.models.quantization as models

# You will need the number of filters in the `fc` for future use.
# Here the size of each output sample is set to 2.
# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
model_fe = models.resnet18(pretrained=True, progress=True, quantize=True)
num_ftrs = model_fe.fc.in_features

在此时你需要修改预训练模型。该模型在开始和结束处包含量化/反量化块。然而，因为你只会使用特征提取器，所以反量化层必须移动到线性层（头部）之前。最简单的方法是将模型包装在nn.Sequential模块中。

第一步是将ResNet模型中的特征提取器隔离出来。尽管在这个示例中你被要求使用除fc之外的所有层作为特征提取器，但实际上你可以根据需要选择任意部分。这在你想要替换某些卷积层的情况下会非常有用。

注意

当将量化模型中的特征提取器与其他部分分离时，你必须手动在你希望保留量化部分的开始和结束位置放置量化器/反量化器。

下面的函数创建了一个带有自定义头部的模型。

from torch import nn

def create_combined_model(model_fe):
  # Step 1. Isolate the feature extractor.
  model_fe_features = nn.Sequential(
    model_fe.quant,  # Quantize the input
    model_fe.conv1,
    model_fe.bn1,
    model_fe.relu,
    model_fe.maxpool,
    model_fe.layer1,
    model_fe.layer2,
    model_fe.layer3,
    model_fe.layer4,
    model_fe.avgpool,
    model_fe.dequant,  # Dequantize the output
  )

  # Step 2. Create a new "head"
  new_head = nn.Sequential(
    nn.Dropout(p=0.5),
    nn.Linear(num_ftrs, 2),
  )

  # Step 3. Combine, and don't forget the quant stubs.
  new_model = nn.Sequential(
    model_fe_features,
    nn.Flatten(1),
    new_head,
  )
  return new_model

警告

目前量化模型只能在 CPU 上运行。 然而，可以将模型中未量化的部分发送到 GPU 上。

import torch.optim as optim
new_model = create_combined_model(model_fe)
new_model = new_model.to('cpu')

criterion = nn.CrossEntropyLoss()

# Note that we are only training the head.
optimizer_ft = optim.SGD(new_model.parameters(), lr=0.01, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

训练和评估

这一步在 CPU 上大约需要 15-25 分钟。因为量化模型只能在 CPU 上运行，因此无法在 GPU 上进行训练。

new_model = train_model(new_model, criterion, optimizer_ft, exp_lr_scheduler,
                        num_epochs=25, device='cpu')

visualize_model(new_model)
plt.tight_layout()

第2部分. 微调可量化模型

在这一部分，我们对用于迁移学习的特征提取器进行微调，并对特征提取器进行量化。请注意，在第1部分和第2部分中，特征提取器都会被量化。不同之处在于，第1部分中我们使用的是预训练的量化模型。而在这一部分中，我们在感兴趣的数据集上对模型进行微调后，再创建量化特征提取器，因此这是一种在实现迁移学习时获得更好准确率的同时，还能享受量化优势的方法。需要注意的是，在我们的具体示例中，训练集非常小（只有120张图片），因此对整个模型进行微调的优势并不明显。然而，此处展示的流程对于更大的数据集进行迁移学习时将有助于提高准确率。

预训练的特征提取器必须可量化。 为确保其可量化，请执行以下步骤：

1. Fuse (Conv, BN, ReLU), (Conv, BN), and (Conv, ReLU) using torch.quantization.fuse_modules.

2. Connect the feature extractor with a custom head. This requires dequantizing the output of the feature extractor.

3. Insert fake-quantization modules at appropriate locations in the feature extractor to mimic quantization during training.

对于步骤 (1)，我们使用来自 torchvision/models/quantization 的模型， 这些模型具有一个成员方法 fuse_model。此函数融合所有 conv， bn 和 relu 模块。对于自定义模型，这将需要手动调用 torch.quantization.fuse_modules API，并传入要融合的模块列表。

步骤 (2) 由上一节中使用的 create_combined_model 函数执行。

步骤 (3) 通过使用 torch.quantization.prepare_qat 实现，该操作 插入了假量化模块。

如步骤 (4) 所述，您可以开始“微调”模型，然后将其转换为完全量化版本（步骤 5）。

要将微调后的模型转换为量化模型，您可以调用 torch.quantization.convert 函数（在我们的情况下，仅 特征提取器被量化）。

注意

由于随机初始化，您的结果可能与本教程中展示的结果有所不同。

# notice `quantize=False`
model = models.resnet18(pretrained=True, progress=True, quantize=False)
num_ftrs = model.fc.in_features

# Step 1
model.train()
model.fuse_model()
# Step 2
model_ft = create_combined_model(model)
model_ft[0].qconfig = torch.quantization.default_qat_qconfig  # Use default QAT configuration
# Step 3
model_ft = torch.quantization.prepare_qat(model_ft, inplace=True)

微调模型

在当前的教程中，整个模型进行了微调。一般来说，这会导致更高的准确性。然而，由于这里使用的训练集较小，我们最终会过度拟合到训练集上。

第4步. 微调模型

for param in model_ft.parameters():
  param.requires_grad = True

model_ft.to(device)  # We can fine-tune on GPU if available

criterion = nn.CrossEntropyLoss()

# Note that we are training everything, so the learning rate is lower
# Notice the smaller learning rate
optimizer_ft = optim.SGD(model_ft.parameters(), lr=1e-3, momentum=0.9, weight_decay=0.1)

# Decay LR by a factor of 0.3 every several epochs
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.3)

model_ft_tuned = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                             num_epochs=25, device=device)

步骤 5. 转换为量化模型

from torch.quantization import convert
model_ft_tuned.cpu()

model_quantized_and_trained = convert(model_ft_tuned, inplace=False)

让我们看看量化模型在几张图像上的表现

visualize_model(model_quantized_and_trained)

plt.ioff()
plt.tight_layout()
plt.show()