(beta) 用於電腦視覺的量化遷移學習教程

建立於：2019 年 12 月 06 日 | 最後更新：2021 年 7 月 27 日 | 最後驗證：2024 年 11 月 05 日

提示

為了充分利用本教程，我們建議使用此Colab 版本。 這樣您就可以試驗下面介紹的資訊。

作者：Zafar Takhirov

審閱者：Raghuraman Krishnamoorthi

編輯者：Jessica Lin

本教程基於 PyTorch 遷移學習 原始教程，由 Sasank Chilamkurthy 撰寫。

遷移學習是指利用預訓練模型來應用於不同資料集的技术。遷移學習主要有兩種使用方式

1. ConvNet 作為固定的特徵提取器：在此，您可以“凍結”網路中所有參數的權重，除了最後幾層（又名“頭部”，通常是完全連接的層）。這些最後幾層將被替換為以隨機權重初始化的新層，並且僅訓練這些層。

2. 微調 ConvNet：模型不是使用隨機初始化，而是使用預訓練的網路初始化，然後像往常一樣繼續訓練，但使用不同的資料集。如果輸出數量不同，通常也會在網路中替換頭部（或其一部分）。在這種方法中，通常將學習速率設定為較小的數字。這樣做的原因是網路已經過訓練，只需要進行微小的更改即可將其“微調”到新的資料集。

您還可以組合以上兩種方法：首先，您可以凍結特徵提取器，然後訓練頭部。之後，您可以解凍特徵提取器（或其一部分），將學習速率設定為較小的數值，然後繼續訓練。

在本部分中，您將使用第一種方法 – 使用量化模型提取特徵。

第 0 部分。先決條件

在深入研究遷移學習之前，讓我們先回顧一下“先決條件”，例如安裝和資料載入/視覺化。

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

plt.ion()

安裝 Nightly Build

因為您將使用 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)

第一部分：基於量化特徵提取器訓練自定義分類器

在本節中，您將使用「凍結」的量化特徵提取器，並在其之上訓練自定義分類器 head。與浮點模型不同，您不需要為量化模型設定 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

此時，您需要修改預訓練模型。該模型在開始和結束時具有量化/反量化區塊。但是，由於您只會使用特徵提取器，因此反量化層必須移動到線性層（head）之前。最簡單的方法是將模型包裝在 nn.Sequential 模組中。

第一步是將 ResNet 模型中的特徵提取器隔離出來。雖然在本例中，您的任務是使用除了 fc 之外的所有層作為特徵提取器，但在現實中，您可以根據需要提取任意多的部分。如果您想替換一些卷積層，這會很有用。

注意

從量化模型的其餘部分分離特徵提取器時，您必須手動將量化器/反量化器放置在您要保持量化的部分的開始和結束位置。

下面的函數會建立一個帶有自定義 head 的模型。

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

第二部分：微調可量化模型

在本部分中，我們微調用於遷移學習的特徵提取器，並量化特徵提取器。請注意，在第 1 部分和第 2 部分中，特徵提取器都是量化的。不同之處在於，在第 1 部分中，我們使用預訓練的量化模型。在本部分中，我們在對感興趣的資料集進行微調後，建立一個量化的特徵提取器，因此這是一種在使用遷移學習的同時獲得更好準確性的方法，同時兼顧量化的優點。請注意，在我們特定的示例中，訓練集非常小（120 張圖片），因此微調整個模型的好處並不顯著。但是，此處顯示的程序將提高使用較大資料集進行遷移學習的準確性。

預訓練的特徵提取器必須是可量化的。為了確保它是可量化的，請執行以下步驟

1. 使用 torch.quantization.fuse_modules 融合 (Conv, BN, ReLU)、(Conv, BN) 和 (Conv, ReLU)。

2. 將特徵提取器與自定義 head 連接。這需要反量化特徵提取器的輸出。

3. 在特徵提取器中的適當位置插入 fake-quantization 模組，以在訓練期間模擬量化。

對於步驟 (1)，我們使用 torchvision/models/quantization 中的模型，這些模型具有成員方法 fuse_model。此函數會融合所有 conv、bn 和 relu 模組。對於自定義模型，這需要使用要融合的模組清單手動呼叫 torch.quantization.fuse_modules API。

步驟 (2) 由上一節中使用的 create_combined_model 函數執行。

步驟 (3) 通過使用 torch.quantization.prepare_qat 來實現，該函數會插入 fake-quantization 模組。

作為步驟 (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()