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

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

小费

为了充分利用本教程，我们建议使用此 Colab 版本。 这将使您可以尝试以下信息。

作者： Zafar Takhirov

由审核： Raghuraman Krishnamoorthi

编辑：Jessica Lin

本教程以 Sasank Chilamkurthy 编写的原始 PyTorch 迁移学习教程为基础。

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

  1. 作为固定特征提取器的 ConvNet：在这里，您“冻结”网络中所有参数的权重，除了最后几层（又称“头部”，通常是全连接层）。 将这些最后一层替换为使用随机权重初始化的新层，并且仅训练这些层。
  2. ConvNet 的微调：使用随机训练的网络初始化模型，而不是随机初始化，然后像往常一样使用不同的数据集进行训练。 通常，如果输出数量不同，则在网络中也会更换头部（或头部的一部分）。 这种方法通常将学习率设置为较小的值。 这样做是因为已经对网络进行了训练，并且只需进行较小的更改即可将其“微调”到新的数据集。

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

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

第 0 部分，先决条件

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

  1. # Imports
  2. import copy
  3. import matplotlib.pyplot as plt
  4. import numpy as np
  5. import os
  6. import time
  8. plt.ion()

安装每夜构建

因为您将使用 PyTorch 的 Beta 部分，所以建议安装最新版本的torch和torchvision。 您可以在这里找到有关本地安装的最新说明。 例如，要在没有 GPU 支持的情况下进行安装：

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

加载数据

注意

本部分与原始的迁移学习教程相同。

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

您今天要解决的问题是从图像中对蚂蚁和蜜蜂进行分类。 该数据集包含约 120 张针对蚂蚁和蜜蜂的训练图像。 每个类别有 75 个验证图像。 可以认为这是一个很小的数据集。 但是，由于我们正在使用迁移学习，因此我们应该能够很好地进行概括。

此数据集是 imagenet 的很小子集。

注意

从此处下载数据，并将其提取到data目录。

  1. import torch
  2. from torchvision import transforms, datasets
  4. # Data augmentation and normalization for training
  5. # Just normalization for validation
  6. data_transforms = {
  7.     'train': transforms.Compose([
  8.         transforms.Resize(224),
  9.         transforms.RandomCrop(224),
  10.         transforms.RandomHorizontalFlip(),
  11.         transforms.ToTensor(),
  12.         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  13.     ]),
  14.     'val': transforms.Compose([
  15.         transforms.Resize(224),
  16.         transforms.CenterCrop(224),
  17.         transforms.ToTensor(),
  18.         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  19.     ]),
  20. }
  22. data_dir = 'data/hymenoptera_data'
  23. image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
  24.                                           data_transforms[x])
  25.                   for x in ['train', 'val']}
  26. dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=16,
  27.                                               shuffle=True, num_workers=8)
  28.               for x in ['train', 'val']}
  29. dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
  30. class_names = image_datasets['train'].classes
  32. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

可视化一些图像

让我们可视化一些训练图像，以了解数据扩充。

  1. import torchvision
  3. def imshow(inp, title=None, ax=None, figsize=(5, 5)):
  4.   """Imshow for Tensor."""
  5.   inp = inp.numpy().transpose((1, 2, 0))
  6.   mean = np.array([0.485, 0.456, 0.406])
  7.   std = np.array([0.229, 0.224, 0.225])
  8.   inp = std * inp + mean
  9.   inp = np.clip(inp, 0, 1)
  10.   if ax is None:
  11.     fig, ax = plt.subplots(1, figsize=figsize)
  12.   ax.imshow(inp)
  13.   ax.set_xticks([])
  14.   ax.set_yticks([])
  15.   if title is not None:
  16.     ax.set_title(title)
  18. # Get a batch of training data
  19. inputs, classes = next(iter(dataloaders['train']))
  21. # Make a grid from batch
  22. out = torchvision.utils.make_grid(inputs, nrow=4)
  24. fig, ax = plt.subplots(1, figsize=(10, 10))
  25. imshow(out, title=[class_names[x] for x in classes], ax=ax)

模型训练的支持函数

以下是模型训练的通用函数。 此函数也：

  • 安排学习率
  • 保存最佳模型
  1. def train_model(model, criterion, optimizer, scheduler, num_epochs=25, device='cpu'):
  2.   """
  3.   Support function for model training.
  5.   Args:
  6.     model: Model to be trained
  7.     criterion: Optimization criterion (loss)
  8.     optimizer: Optimizer to use for training
  9.     scheduler: Instance of ``torch.optim.lr_scheduler``
  10.     num_epochs: Number of epochs
  11.     device: Device to run the training on. Must be 'cpu' or 'cuda'
  12.   """
  13.   since = time.time()
  15.   best_model_wts = copy.deepcopy(model.state_dict())
  16.   best_acc = 0.0
  18.   for epoch in range(num_epochs):
  19.     print('Epoch {}/{}'.format(epoch, num_epochs - 1))
  20.     print('-' * 10)
  22.     # Each epoch has a training and validation phase
  23.     for phase in ['train', 'val']:
  24.       if phase == 'train':
  25.         model.train()  # Set model to training mode
  26.       else:
  27.         model.eval()   # Set model to evaluate mode
  29.       running_loss = 0.0
  30.       running_corrects = 0
  32.       # Iterate over data.
  33.       for inputs, labels in dataloaders[phase]:
  34.         inputs = inputs.to(device)
  35.         labels = labels.to(device)
  37.         # zero the parameter gradients
  38.         optimizer.zero_grad()
  40.         # forward
  41.         # track history if only in train
  42.         with torch.set_grad_enabled(phase == 'train'):
  43.           outputs = model(inputs)
  44.           _, preds = torch.max(outputs, 1)
  45.           loss = criterion(outputs, labels)
  47.           # backward + optimize only if in training phase
  48.           if phase == 'train':
  49.             loss.backward()
  50.             optimizer.step()
  52.         # statistics
  53.         running_loss += loss.item() * inputs.size(0)
  54.         running_corrects += torch.sum(preds == labels.data)
  55.       if phase == 'train':
  56.         scheduler.step()
  58.       epoch_loss = running_loss / dataset_sizes[phase]
  59.       epoch_acc = running_corrects.double() / dataset_sizes[phase]
  61.       print('{} Loss: {:.4f} Acc: {:.4f}'.format(
  62.         phase, epoch_loss, epoch_acc))
  64.       # deep copy the model
  65.       if phase == 'val' and epoch_acc > best_acc:
  66.         best_acc = epoch_acc
  67.         best_model_wts = copy.deepcopy(model.state_dict())
  69.     print()
  71.   time_elapsed = time.time() - since
  72.   print('Training complete in {:.0f}m {:.0f}s'.format(
  73.     time_elapsed // 60, time_elapsed % 60))
  74.   print('Best val Acc: {:4f}'.format(best_acc))
  76.   # load best model weights
  77.   model.load_state_dict(best_model_wts)
  78.   return model

可视化模型预测的支持函数

通用函数，显示一些图像的预测

  1. def visualize_model(model, rows=3, cols=3):
  2.   was_training = model.training
  3.   model.eval()
  4.   current_row = current_col = 0
  5.   fig, ax = plt.subplots(rows, cols, figsize=(cols*2, rows*2))
  7.   with torch.no_grad():
  8.     for idx, (imgs, lbls) in enumerate(dataloaders['val']):
  9.       imgs = imgs.cpu()
  10.       lbls = lbls.cpu()
  12.       outputs = model(imgs)
  13.       _, preds = torch.max(outputs, 1)
  15.       for jdx in range(imgs.size()[0]):
  16.         imshow(imgs.data[jdx], ax=ax[current_row, current_col])
  17.         ax[current_row, current_col].axis('off')
  18.         ax[current_row, current_col].set_title('predicted: {}'.format(class_names[preds[jdx]]))
  20.         current_col += 1
  21.         if current_col >= cols:
  22.           current_row += 1
  23.           current_col = 0
  24.         if current_row >= rows:
  25.           model.train(mode=was_training)
  26.           return
  27.     model.train(mode=was_training)

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

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

加载预训练的模型：在本练习中，您将使用 ResNet-18 。

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

此时，您需要修改预训练模型。 该模型在开始和结束时都有量化/去量化块。 但是，由于只使用特征提取器，因此反量化层必须在线性层（头部）之前移动。 最简单的方法是将模型包装在nn.Sequential模块中。

第一步是在 ResNet 模型中隔离特征提取器。 尽管在本示例中，您被责成使用fc以外的所有层作为特征提取器，但实际上，您可以根据需要选择任意数量的零件。 如果您也想替换一些卷积层，这将很有用。

注意

将特征提取器与量化模型的其余部分分开时，必须手动将量化器/去量化器放置在要保持量化的部分的开头和结尾。

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

  1. from torch import nn
  3. def create_combined_model(model_fe):
  4.   # Step 1\. Isolate the feature extractor.
  5.   model_fe_features = nn.Sequential(
  6.     model_fe.quant,  # Quantize the input
  7.     model_fe.conv1,
  8.     model_fe.bn1,
  9.     model_fe.relu,
  10.     model_fe.maxpool,
  11.     model_fe.layer1,
  12.     model_fe.layer2,
  13.     model_fe.layer3,
  14.     model_fe.layer4,
  15.     model_fe.avgpool,
  16.     model_fe.dequant,  # Dequantize the output
  17.   )
  19.   # Step 2\. Create a new "head"
  20.   new_head = nn.Sequential(
  21.     nn.Dropout(p=0.5),
  22.     nn.Linear(num_ftrs, 2),
  23.   )
  25.   # Step 3\. Combine, and don't forget the quant stubs.
  26.   new_model = nn.Sequential(
  27.     model_fe_features,
  28.     nn.Flatten(1),
  29.     new_head,
  30.   )
  31.   return new_model

警告

当前，量化模型只能在 CPU 上运行。 但是，可以将模型的未量化部分发送到 GPU。

  1. import torch.optim as optim
  2. new_model = create_combined_model(model_fe)
  3. new_model = new_model.to('cpu')
  5. criterion = nn.CrossEntropyLoss()
  7. # Note that we are only training the head.
  8. optimizer_ft = optim.SGD(new_model.parameters(), lr=0.01, momentum=0.9)
  10. # Decay LR by a factor of 0.1 every 7 epochs
  11. exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

训练和评估

此步骤在 CPU 上大约需要 15-25 分钟。 由于量化模型只能在 CPU 上运行，因此您不能在 GPU 上运行训练。

  1. new_model = train_model(new_model, criterion, optimizer_ft, exp_lr_scheduler,
  2.                         num_epochs=25, device='cpu')
  4. visualize_model(new_model)
  5. plt.tight_layout()

第 2 部分，微调量化模型

在这一部分中，我们将微调用于迁移学习的特征提取器，并对特征提取器进行量化。 请注意，在第 1 部分和第 2 部分中，特征提取器都是量化的。 不同之处在于，在第 1 部分中，我们使用了预训练的量化模型。 在这一部分中，我们将在对感兴趣的数据集进行微调之后创建一个量化的特征提取器，因此这是一种在具有量化优势的同时通过迁移学习获得更好的准确率的方法。 请注意，在我们的特定示例中，训练集非常小（120 张图像），因此微调整个模型的好处并不明显。 但是，此处显示的过程将提高使用较大数据集进行传递学习的准确率。

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

  1. 使用 torch.quantization.fuse_modules 熔断 (Conv, BN, ReLU) ， (Conv, BN) 和 (Conv, ReLU) 。
  2. 将特征提取器与自定义头部连接。 这需要对特征提取器的输出进行反量化。
  3. 在特征提取器中的适当位置插入伪量化模块，以模拟训练期间的量化。

对于步骤（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函数（在我们的情况下，仅对特征提取器进行量化）。

注意

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

  1. # notice quantize=False model = models.resnet18(pretrained=True, progress=True, quantize=False) num_ftrs = model.fc.in_features
  3. # 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.微调模型

  1. for param in model_ft.parameters():
  2.   param.requires_grad = True
  4. model_ft.to(device)  # We can fine-tune on GPU if available
  6. criterion = nn.CrossEntropyLoss()
  8. # Note that we are training everything, so the learning rate is lower
  9. # Notice the smaller learning rate
  10. optimizer_ft = optim.SGD(model_ft.parameters(), lr=1e-3, momentum=0.9, weight_decay=0.1)
  12. # Decay LR by a factor of 0.3 every several epochs
  13. exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.3)
  15. model_ft_tuned = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
  16.                              num_epochs=25, device=device)

步骤 5.转换为量化模型

  1. from torch.quantization import convert
  2. model_ft_tuned.cpu()
  4. model_quantized_and_trained = convert(model_ft_tuned, inplace=False)

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

  1. visualize_model(model_quantized_and_trained)
  3. plt.ioff()
  4. plt.tight_layout()
  5. plt.show()