PyTorch 2.2 中文官方教程（十五）（1）-阿里云开发者社区

（beta）计算机视觉的量化迁移学习教程

原文：pytorch.org/tutorials/intermediate/quantized_transfer_learning_tutorial.html

译者：飞龙

协议：CC BY-NC-SA 4.0

提示

为了充分利用本教程，我们建议使用这个Colab 版本。这将允许您尝试下面提供的信息。

作者：Zafar Takhirov

审阅者：Raghuraman Krishnamoorthi

编辑：Jessica Lin

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

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

将 ConvNet 作为固定特征提取器：在这里，您会“冻结”网络中除最后几层（通常是完全连接的层，也称为“头部”）之外的所有参数的权重。这些最后的层将被新的层替换，并用随机权重初始化，只有这些层会被训练。
微调 ConvNet：不是随机初始化，而是使用预训练网络初始化模型，然后训练过程与通常情况下不同数据集的训练相同。通常还会替换网络中的头部（或其中的一部分），以适应不同数量的输出。在这种方法中，通常将学习率设置为较小的值。这是因为网络已经训练过，只需要对其进行“微调”以适应新数据集。

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

在本部分中，您将使用第一种方法——使用量化模型提取特征。

第 0 部分。先决条件

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

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

安装夜间版本

由于您将使用 PyTorch 的 beta 部分，建议安装最新版本的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()