PyTorch 中使用 Eager 模式的静态量化（beta）

原文：https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html#

作者： Raghuraman Krishnamoorthi

编辑：Seth Weidman

本教程说明了如何进行训练后的静态量化，并说明了两种更先进的技术-每通道量化和量化感知训练-可以进一步提高模型的准确率。 请注意，目前仅支持 CPU 量化，因此在本教程中我们将不使用 GPU/CUDA。

在本教程结束时，您将看到 PyTorch 中的量化如何导致模型大小显着减小同时提高速度。 此外，您将在此处看到如何轻松应用显示的一些高级量化技术，以使您的量化模型受到的准确率影响要小得多。

警告：我们使用了许多其他 PyTorch 仓库中的样板代码，例如，定义了MobileNetV2模型架构，定义了数据加载器等等。 我们当然鼓励您阅读它； 但是如果要使用量化功能，请随时跳到“ 4。 训练后静态量化”部分。

我们将从进行必要的导入开始：

  1. import numpy as np
  2. import torch
  3. import torch.nn as nn
  4. import torchvision
  5. from torch.utils.data import DataLoader
  6. from torchvision import datasets
  7. import torchvision.transforms as transforms
  8. import os
  9. import time
  10. import sys
  11. import torch.quantization
  13. # # Setup warnings
  14. import warnings
  15. warnings.filterwarnings(
  16.     action='ignore',
  17.     category=DeprecationWarning,
  18.     module=r'.*'
  19. )
  20. warnings.filterwarnings(
  21.     action='default',
  22.     module=r'torch.quantization'
  23. )
  25. # Specify random seed for repeatable results
  26. torch.manual_seed(191009)

1.模型架构

我们首先定义 MobileNetV2 模型架构，并进行了一些值得注意的修改以实现量化：

  • 用nn.quantized.FloatFunctional代替添加
  • 在网络的开头和结尾处插入QuantStub和DeQuantStub。
  • 用 ReLU 替换 ReLU6

注意：此代码取自此处。

  1. from torch.quantization import QuantStub, DeQuantStub
  3. def _make_divisible(v, divisor, min_value=None):
  4.     """
  5.     This function is taken from the original tf repo.
  6.     It ensures that all layers have a channel number that is divisible by 8
  7.     It can be seen here:
  8.     https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
  9.     :param v:
  10.     :param divisor:
  11.     :param min_value:
  12.     :return:
  13.     """
  14.     if min_value is None:
  15.         min_value = divisor
  16.     new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
  17.     # Make sure that round down does not go down by more than 10%.
  18.     if new_v < 0.9 * v:
  19.         new_v += divisor
  20.     return new_v
  22. class ConvBNReLU(nn.Sequential):
  23.     def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
  24.         padding = (kernel_size - 1) // 2
  25.         super(ConvBNReLU, self).__init__(
  26.             nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
  27.             nn.BatchNorm2d(out_planes, momentum=0.1),
  28.             # Replace with ReLU
  29.             nn.ReLU(inplace=False)
  30.         )
  32. class InvertedResidual(nn.Module):
  33.     def __init__(self, inp, oup, stride, expand_ratio):
  34.         super(InvertedResidual, self).__init__()
  35.         self.stride = stride
  36.         assert stride in [1, 2]
  38.         hidden_dim = int(round(inp * expand_ratio))
  39.         self.use_res_connect = self.stride == 1 and inp == oup
  41.         layers = []
  42.         if expand_ratio != 1:
  43.             # pw
  44.             layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
  45.         layers.extend([
  46.             # dw
  47.             ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
  48.             # pw-linear
  49.             nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
  50.             nn.BatchNorm2d(oup, momentum=0.1),
  51.         ])
  52.         self.conv = nn.Sequential(*layers)
  53.         # Replace torch.add with floatfunctional
  54.         self.skip_add = nn.quantized.FloatFunctional()
  56.     def forward(self, x):
  57.         if self.use_res_connect:
  58.             return self.skip_add.add(x, self.conv(x))
  59.         else:
  60.             return self.conv(x)
  62. class MobileNetV2(nn.Module):
  63.     def __init__(self, num_classes=1000, width_mult=1.0, inverted_residual_setting=None, round_nearest=8):
  64.         """
  65.         MobileNet V2 main class
  67.         Args:
  68.             num_classes (int): Number of classes
  69.             width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
  70.             inverted_residual_setting: Network structure
  71.             round_nearest (int): Round the number of channels in each layer to be a multiple of this number
  72.             Set to 1 to turn off rounding
  73.         """
  74.         super(MobileNetV2, self).__init__()
  75.         block = InvertedResidual
  76.         input_channel = 32
  77.         last_channel = 1280
  79.         if inverted_residual_setting is None:
  80.             inverted_residual_setting = [
  81.                 # t, c, n, s
  82.                 [1, 16, 1, 1],
  83.                 [6, 24, 2, 2],
  84.                 [6, 32, 3, 2],
  85.                 [6, 64, 4, 2],
  86.                 [6, 96, 3, 1],
  87.                 [6, 160, 3, 2],
  88.                 [6, 320, 1, 1],
  89.             ]
  91.         # only check the first element, assuming user knows t,c,n,s are required
  92.         if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
  93.             raise ValueError("inverted_residual_setting should be non-empty "
  94.                              "or a 4-element list, got {}".format(inverted_residual_setting))
  96.         # building first layer
  97.         input_channel = _make_divisible(input_channel * width_mult, round_nearest)
  98.         self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
  99.         features = [ConvBNReLU(3, input_channel, stride=2)]
  100.         # building inverted residual blocks
  101.         for t, c, n, s in inverted_residual_setting:
  102.             output_channel = _make_divisible(c * width_mult, round_nearest)
  103.             for i in range(n):
  104.                 stride = s if i == 0 else 1
  105.                 features.append(block(input_channel, output_channel, stride, expand_ratio=t))
  106.                 input_channel = output_channel
  107.         # building last several layers
  108.         features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
  109.         # make it nn.Sequential
  110.         self.features = nn.Sequential(*features)
  111.         self.quant = QuantStub()
  112.         self.dequant = DeQuantStub()
  113.         # building classifier
  114.         self.classifier = nn.Sequential(
  115.             nn.Dropout(0.2),
  116.             nn.Linear(self.last_channel, num_classes),
  117.         )
  119.         # weight initialization
  120.         for m in self.modules():
  121.             if isinstance(m, nn.Conv2d):
  122.                 nn.init.kaiming_normal_(m.weight, mode='fan_out')
  123.                 if m.bias is not None:
  124.                     nn.init.zeros_(m.bias)
  125.             elif isinstance(m, nn.BatchNorm2d):
  126.                 nn.init.ones_(m.weight)
  127.                 nn.init.zeros_(m.bias)
  128.             elif isinstance(m, nn.Linear):
  129.                 nn.init.normal_(m.weight, 0, 0.01)
  130.                 nn.init.zeros_(m.bias)
  132.     def forward(self, x):
  134.         x = self.quant(x)
  136.         x = self.features(x)
  137.         x = x.mean([2, 3])
  138.         x = self.classifier(x)
  139.         x = self.dequant(x)
  140.         return x
  142.     # Fuse Conv+BN and Conv+BN+Relu modules prior to quantization
  143.     # This operation does not change the numerics
  144.     def fuse_model(self):
  145.         for m in self.modules():
  146.             if type(m) == ConvBNReLU:
  147.                 torch.quantization.fuse_modules(m, ['0', '1', '2'], inplace=True)
  148.             if type(m) == InvertedResidual:
  149.                 for idx in range(len(m.conv)):
  150.                     if type(m.conv[idx]) == nn.Conv2d:
  151.                         torch.quantization.fuse_modules(m.conv, [str(idx), str(idx + 1)], inplace=True)

2.辅助函数

接下来，我们定义一些助手函数以帮助模型评估。 这些主要来自这里。

  1. class AverageMeter(object):
  2.     """Computes and stores the average and current value"""
  3.     def __init__(self, name, fmt=':f'):
  4.         self.name = name
  5.         self.fmt = fmt
  6.         self.reset()
  8.     def reset(self):
  9.         self.val = 0
  10.         self.avg = 0
  11.         self.sum = 0
  12.         self.count = 0
  14.     def update(self, val, n=1):
  15.         self.val = val
  16.         self.sum += val * n
  17.         self.count += n
  18.         self.avg = self.sum / self.count
  20.     def __str__(self):
  21.         fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
  22.         return fmtstr.format(**self.__dict__)
  24. def accuracy(output, target, topk=(1,)):
  25.     """Computes the accuracy over the k top predictions for the specified values of k"""
  26.     with torch.no_grad():
  27.         maxk = max(topk)
  28.         batch_size = target.size(0)
  30.         _, pred = output.topk(maxk, 1, True, True)
  31.         pred = pred.t()
  32.         correct = pred.eq(target.view(1, -1).expand_as(pred))
  34.         res = []
  35.         for k in topk:
  36.             correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
  37.             res.append(correct_k.mul_(100.0 / batch_size))
  38.         return res
  40. def evaluate(model, criterion, data_loader, neval_batches):
  41.     model.eval()
  42.     top1 = AverageMeter('Acc@1', ':6.2f')
  43.     top5 = AverageMeter('Acc@5', ':6.2f')
  44.     cnt = 0
  45.     with torch.no_grad():
  46.         for image, target in data_loader:
  47.             output = model(image)
  48.             loss = criterion(output, target)
  49.             cnt += 1
  50.             acc1, acc5 = accuracy(output, target, topk=(1, 5))
  51.             print('.', end = '')
  52.             top1.update(acc1[0], image.size(0))
  53.             top5.update(acc5[0], image.size(0))
  54.             if cnt >= neval_batches:
  55.                  return top1, top5
  57.     return top1, top5
  59. def load_model(model_file):
  60.     model = MobileNetV2()
  61.     state_dict = torch.load(model_file)
  62.     model.load_state_dict(state_dict)
  63.     model.to('cpu')
  64.     return model
  66. def print_size_of_model(model):
  67.     torch.save(model.state_dict(), "temp.p")
  68.     print('Size (MB):', os.path.getsize("temp.p")/1e6)
  69.     os.remove('temp.p')

3.定义数据集和数据加载器

作为最后的主要设置步骤，我们为训练和测试集定义了数据加载器。

ImageNet 数据

我们为本教程创建的特定数据集仅包含来自 ImageNet 数据的 1000 张图像，每个类别都有一张（此数据集的大小刚好超过 250 MB，可以相对轻松地下载）。 此自定义数据集的 URL 为：

  1. https://s3.amazonaws.com/pytorch-tutorial-assets/imagenet_1k.zip

要使用 Python 在本地下载此数据，可以使用：

  1. import requests
  3. url = 'https://s3.amazonaws.com/pytorch-tutorial-assets/imagenet_1k.zip`
  4. filename = '~/Downloads/imagenet_1k_data.zip'
  6. r = requests.get(url)
  8. with open(filename, 'wb') as f:
  9.     f.write(r.content)

为了运行本教程，我们下载了这些数据，并使用 Makefile 中的这些行将其移到正确的位置。

另一方面，要使用整个 ImageNet 数据集运行本教程中的代码，可以在此之后使用torchvision下载数据。 例如，要下载训练集并对其进行一些标准转换，可以使用：

  1. import torchvision
  2. import torchvision.transforms as transforms
  4. imagenet_dataset = torchvision.datasets.ImageNet(
  5.     '~/.data/imagenet',
  6.     split='train',
  7.     download=True,
  8.     transforms.Compose([
  9.         transforms.RandomResizedCrop(224),
  10.         transforms.RandomHorizontalFlip(),
  11.         transforms.ToTensor(),
  12.         transforms.Normalize(mean=[0.485, 0.456, 0.406],
  13.                              std=[0.229, 0.224, 0.225]),
  14.     ])

下载完数据后，我们在下面显示了一些函数，这些函数定义了将用于读取此数据的数据加载器。 这些函数主要来自此处。

  1. def prepare_data_loaders(data_path):
  3.     traindir = os.path.join(data_path, 'train')
  4.     valdir = os.path.join(data_path, 'val')
  5.     normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
  6.                                      std=[0.229, 0.224, 0.225])
  8.     dataset = torchvision.datasets.ImageFolder(
  9.         traindir,
  10.         transforms.Compose([
  11.             transforms.RandomResizedCrop(224),
  12.             transforms.RandomHorizontalFlip(),
  13.             transforms.ToTensor(),
  14.             normalize,
  15.         ]))
  17.     dataset_test = torchvision.datasets.ImageFolder(
  18.         valdir,
  19.         transforms.Compose([
  20.             transforms.Resize(256),
  21.             transforms.CenterCrop(224),
  22.             transforms.ToTensor(),
  23.             normalize,
  24.         ]))
  26.     train_sampler = torch.utils.data.RandomSampler(dataset)
  27.     test_sampler = torch.utils.data.SequentialSampler(dataset_test)
  29.     data_loader = torch.utils.data.DataLoader(
  30.         dataset, batch_size=train_batch_size,
  31.         sampler=train_sampler)
  33.     data_loader_test = torch.utils.data.DataLoader(
  34.         dataset_test, batch_size=eval_batch_size,
  35.         sampler=test_sampler)
  37.     return data_loader, data_loader_test

接下来，我们将加载经过预先​​训练的 MobileNetV2 模型。 我们在这里提供用于从torchvision中下载数据的 URL。

  1. data_path = 'data/imagenet_1k'
  2. saved_model_dir = 'data/'
  3. float_model_file = 'mobilenet_pretrained_float.pth'
  4. scripted_float_model_file = 'mobilenet_quantization_scripted.pth'
  5. scripted_quantized_model_file = 'mobilenet_quantization_scripted_quantized.pth'
  7. train_batch_size = 30
  8. eval_batch_size = 30
  10. data_loader, data_loader_test = prepare_data_loaders(data_path)
  11. criterion = nn.CrossEntropyLoss()
  12. float_model = load_model(saved_model_dir + float_model_file).to('cpu')

接下来，我们将“融合模块”； 通过节省内存访问量，这可以使模型更快，同时还可以提高数值精度。 尽管这可以用于任何模型，但在量化模型中尤为常见。

  1. print('\n Inverted Residual Block: Before fusion \n\n', float_model.features[1].conv)
  2. float_model.eval()
  4. # Fuses modules
  5. float_model.fuse_model()
  7. # Note fusion of Conv+BN+Relu and Conv+Relu
  8. print('\n Inverted Residual Block: After fusion\n\n',float_model.features[1].conv)

出：

  1. Inverted Residual Block: Before fusion
  3.  Sequential(
  4.   (0): ConvBNReLU(
  5.     (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)
  6.     (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  7.     (2): ReLU()
  8.   )
  9.   (1): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
  10.   (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  11. )
  13.  Inverted Residual Block: After fusion
  15.  Sequential(
  16.   (0): ConvBNReLU(
  17.     (0): ConvReLU2d(
  18.       (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32)
  19.       (1): ReLU()
  20.     )
  21.     (1): Identity()
  22.     (2): Identity()
  23.   )
  24.   (1): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1))
  25.   (2): Identity()
  26. )

最后，为了获得“基准”精度，让我们看看带有融合模块的未量化模型的精度

  1. num_eval_batches = 10
  3. print("Size of baseline model")
  4. print_size_of_model(float_model)
  6. top1, top5 = evaluate(float_model, criterion, data_loader_test, neval_batches=num_eval_batches)
  7. print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))
  8. torch.jit.save(torch.jit.script(float_model), saved_model_dir + scripted_float_model_file)

出：

  1. Size of baseline model
  2. Size (MB): 13.999657
  3. ..........Evaluation accuracy on 300 images, 77.67

我们看到 300 张图像的准确率达到 78%，这是 ImageNet 的坚实基础，特别是考虑到我们的模型只有 14.0 MB。

这将是我们比较的基准。 接下来，让我们尝试不同的量化方法

4.训练后的静态量化

训练后的静态量化不仅涉及像动态量化中那样将权重从float转换为int，而且还执行额外的步骤，即首先通过网络馈送一批数据并计算不同激活的结果分布（具体来说，这通过在记录此数据的不同点插入观察者模块来完成）。 然后使用这些分布来确定在推理时如何具体量化不同的激活（一种简单的技术是将整个激活范围简单地划分为 256 个级别，但我们也支持更复杂的方法）。 重要的是，此附加步骤使我们能够在操作之间传递量化值，而不是在每次操作之间将这些值转换为浮点数，然后再转换为整数，从而显着提高了速度。

  1. num_calibration_batches = 10
  3. myModel = load_model(saved_model_dir + float_model_file).to('cpu')
  4. myModel.eval()
  6. # Fuse Conv, bn and relu
  7. myModel.fuse_model()
  9. # Specify quantization configuration
  10. # Start with simple min/max range estimation and per-tensor quantization of weights
  11. myModel.qconfig = torch.quantization.default_qconfig
  12. print(myModel.qconfig)
  13. torch.quantization.prepare(myModel, inplace=True)
  15. # Calibrate first
  16. print('Post Training Quantization Prepare: Inserting Observers')
  17. print('\n Inverted Residual Block:After observer insertion \n\n', myModel.features[1].conv)
  19. # Calibrate with the training set
  20. evaluate(myModel, criterion, data_loader, neval_batches=num_calibration_batches)
  21. print('Post Training Quantization: Calibration done')
  23. # Convert to quantized model
  24. torch.quantization.convert(myModel, inplace=True)
  25. print('Post Training Quantization: Convert done')
  26. print('\n Inverted Residual Block: After fusion and quantization, note fused modules: \n\n',myModel.features[1].conv)
  28. print("Size of model after quantization")
  29. print_size_of_model(myModel)
  31. top1, top5 = evaluate(myModel, criterion, data_loader_test, neval_batches=num_eval_batches)
  32. print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))

出：

  1. QConfig(activation=functools.partial(<class 'torch.quantization.observer.MinMaxObserver'>, reduce_range=True), weight=functools.partial(<class 'torch.quantization.observer.MinMaxObserver'>, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric))
  2. Post Training Quantization Prepare: Inserting Observers
  4.  Inverted Residual Block:After observer insertion
  6.  Sequential(
  7.   (0): ConvBNReLU(
  8.     (0): ConvReLU2d(
  9.       (0): Conv2d(
  10.         32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32
  11.         (activation_post_process): MinMaxObserver(min_val=inf, max_val=-inf)
  12.       )
  13.       (1): ReLU(
  14.         (activation_post_process): MinMaxObserver(min_val=inf, max_val=-inf)
  15.       )
  16.     )
  17.     (1): Identity()
  18.     (2): Identity()
  19.   )
  20.   (1): Conv2d(
  21.     32, 16, kernel_size=(1, 1), stride=(1, 1)
  22.     (activation_post_process): MinMaxObserver(min_val=inf, max_val=-inf)
  23.   )
  24.   (2): Identity()
  25. )
  26. ..........Post Training Quantization: Calibration done
  27. Post Training Quantization: Convert done
  29.  Inverted Residual Block: After fusion and quantization, note fused modules:
  31.  Sequential(
  32.   (0): ConvBNReLU(
  33.     (0): QuantizedConvReLU2d(32, 32, kernel_size=(3, 3), stride=(1, 1), scale=0.1516050398349762, zero_point=0, padding=(1, 1), groups=32)
  34.     (1): Identity()
  35.     (2): Identity()
  36.   )
  37.   (1): QuantizedConv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), scale=0.17719413340091705, zero_point=63)
  38.   (2): Identity()
  39. )
  40. Size of model after quantization
  41. Size (MB): 3.631847
  42. ..........Evaluation accuracy on 300 images, 66.67

对于这个量化模型，我们发现在这 300 张相同的图像上，准确率仅低至约 62%。 不过，我们确实将模型的大小减小到了 3.6 MB 以下，几乎减少了 4 倍。

此外，我们可以通过使用不同的量化配置来显着提高准确率。 我们使用推荐的配置对 x86 架构进行量化，重复相同的练习。 此配置执行以下操作：

  • 量化每个通道的权重
  • 使用直方图观察器，该直方图观察器收集激活的直方图，然后以最佳方式选择量化参数。
  1. per_channel_quantized_model = load_model(saved_model_dir + float_model_file)
  2. per_channel_quantized_model.eval()
  3. per_channel_quantized_model.fuse_model()
  4. per_channel_quantized_model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
  5. print(per_channel_quantized_model.qconfig)
  7. torch.quantization.prepare(per_channel_quantized_model, inplace=True)
  8. evaluate(per_channel_quantized_model,criterion, data_loader, num_calibration_batches)
  9. torch.quantization.convert(per_channel_quantized_model, inplace=True)
  10. top1, top5 = evaluate(per_channel_quantized_model, criterion, data_loader_test, neval_batches=num_eval_batches)
  11. print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))
  12. torch.jit.save(torch.jit.script(per_channel_quantized_model), saved_model_dir + scripted_quantized_model_file)

出：

  1. QConfig(activation=functools.partial(<class 'torch.quantization.observer.HistogramObserver'>, reduce_range=True), weight=functools.partial(<class 'torch.quantization.observer.PerChannelMinMaxObserver'>, dtype=torch.qint8, qscheme=torch.per_channel_symmetric))
  2. ....................Evaluation accuracy on 300 images, 74.67

仅更改这种量化配置方法，就可以将准确率提高到 76% 以上！ 尽管如此，这仍比上述 78% 的基准差 1-2%。 因此，让我们尝试量化意识的训练。

5.量化感知的训练

量化感知的训练（QAT）是通常导致最高准确率的量化方法。 使用 QAT，在训练的正向和反向过程中，所有权重和激活都被“伪量化”：即，浮点值四舍五入以模拟int8值，但所有计算仍使用浮点数完成。 因此，在训练过程中进行所有权重调整，同时“意识到”该模型将最终被量化的事实。 因此，在量化之后，此方法通常会比动态量化或训练后静态量化产生更高的精度。

实际执行 QAT 的总体工作流程与之前非常相似：

  • 我们可以使用与以前相同的模型：量化感知的训练不需要额外的准备。
  • 我们需要使用qconfig来指定要在权重和激活之后插入哪种伪量化，而不是指定观察者

我们首先定义一个训练函数：

  1. def train_one_epoch(model, criterion, optimizer, data_loader, device, ntrain_batches):
  2.     model.train()
  3.     top1 = AverageMeter('Acc@1', ':6.2f')
  4.     top5 = AverageMeter('Acc@5', ':6.2f')
  5.     avgloss = AverageMeter('Loss', '1.5f')
  7.     cnt = 0
  8.     for image, target in data_loader:
  9.         start_time = time.time()
  10.         print('.', end = '')
  11.         cnt += 1
  12.         image, target = image.to(device), target.to(device)
  13.         output = model(image)
  14.         loss = criterion(output, target)
  15.         optimizer.zero_grad()
  16.         loss.backward()
  17.         optimizer.step()
  18.         acc1, acc5 = accuracy(output, target, topk=(1, 5))
  19.         top1.update(acc1[0], image.size(0))
  20.         top5.update(acc5[0], image.size(0))
  21.         avgloss.update(loss, image.size(0))
  22.         if cnt >= ntrain_batches:
  23.             print('Loss', avgloss.avg)
  25.             print('Training: * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'
  26.                   .format(top1=top1, top5=top5))
  27.             return
  29.     print('Full imagenet train set:  * Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}'
  30.           .format(top1=top1, top5=top5))
  31.     return

我们像以前一样融合模块

  1. qat_model = load_model(saved_model_dir + float_model_file)
  2. qat_model.fuse_model()
  4. optimizer = torch.optim.SGD(qat_model.parameters(), lr = 0.0001)
  5. qat_model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

最后，prepare_qat执行“伪量化”，为量化感知训练准备模型

  1. torch.quantization.prepare_qat(qat_model, inplace=True)
  2. print('Inverted Residual Block: After preparation for QAT, note fake-quantization modules \n',qat_model.features[1].conv)

出：

  1. Inverted Residual Block: After preparation for QAT, note fake-quantization modules
  2.  Sequential(
  3.   (0): ConvBNReLU(
  4.     (0): ConvBnReLU2d(
  5.       32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False
  6.       (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  7.       (weight_fake_quant): FakeQuantize(
  8.         fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=torch.qint8, qscheme=torch.per_channel_symmetric, ch_axis=0,         scale=tensor([1.]), zero_point=tensor([0])
  9.         (activation_post_process): MovingAveragePerChannelMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
  10.       )
  11.       (activation_post_process): FakeQuantize(
  12.         fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1,         scale=tensor([1.]), zero_point=tensor([0])
  13.         (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)
  14.       )
  15.     )
  16.     (1): Identity()
  17.     (2): Identity()
  18.   )
  19.   (1): ConvBn2d(
  20.     32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False
  21.     (bn): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  22.     (weight_fake_quant): FakeQuantize(
  23.       fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=torch.qint8, qscheme=torch.per_channel_symmetric, ch_axis=0,         scale=tensor([1.]), zero_point=tensor([0])
  24.       (activation_post_process): MovingAveragePerChannelMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
  25.     )
  26.     (activation_post_process): FakeQuantize(
  27.       fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1,         scale=tensor([1.]), zero_point=tensor([0])
  28.       (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)
  29.     )
  30.   )
  31.   (2): Identity()
  32. )

高精度训练量化模型需要在推断时对数字进行精确建模。 因此，对于量化感知的训练，我们通过以下方式修改训练循环：

  • 在训练快要结束时切换批量规范以使用运行均值和方差，以更好地匹配推理数字。
  • 我们还冻结了量化器参数（比例和零点），并对权重进行了微调。
  1. num_train_batches = 20
  3. # Train and check accuracy after each epoch
  4. for nepoch in range(8):
  5.     train_one_epoch(qat_model, criterion, optimizer, data_loader, torch.device('cpu'), num_train_batches)
  6.     if nepoch > 3:
  7.         # Freeze quantizer parameters
  8.         qat_model.apply(torch.quantization.disable_observer)
  9.     if nepoch > 2:
  10.         # Freeze batch norm mean and variance estimates
  11.         qat_model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)
  13.     # Check the accuracy after each epoch
  14.     quantized_model = torch.quantization.convert(qat_model.eval(), inplace=False)
  15.     quantized_model.eval()
  16.     top1, top5 = evaluate(quantized_model,criterion, data_loader_test, neval_batches=num_eval_batches)
  17.     print('Epoch %d :Evaluation accuracy on %d images, %2.2f'%(nepoch, num_eval_batches * eval_batch_size, top1.avg))

出：

  1. ....................Loss tensor(2.0747, grad_fn=<DivBackward0>)
  2. Training: * Acc@1 56.167 Acc@5 77.333
  3. ..........Epoch 0 :Evaluation accuracy on 300 images, 77.67
  4. ....................Loss tensor(2.0358, grad_fn=<DivBackward0>)
  5. Training: * Acc@1 54.833 Acc@5 78.500
  6. ..........Epoch 1 :Evaluation accuracy on 300 images, 77.00
  7. ....................Loss tensor(2.0417, grad_fn=<DivBackward0>)
  8. Training: * Acc@1 54.667 Acc@5 77.333
  9. ..........Epoch 2 :Evaluation accuracy on 300 images, 74.67
  10. ....................Loss tensor(1.9055, grad_fn=<DivBackward0>)
  11. Training: * Acc@1 56.833 Acc@5 78.667
  12. ..........Epoch 3 :Evaluation accuracy on 300 images, 76.33
  13. ....................Loss tensor(1.9055, grad_fn=<DivBackward0>)
  14. Training: * Acc@1 58.167 Acc@5 80.000
  15. ..........Epoch 4 :Evaluation accuracy on 300 images, 77.00
  16. ....................Loss tensor(1.7821, grad_fn=<DivBackward0>)
  17. Training: * Acc@1 60.500 Acc@5 82.833
  18. ..........Epoch 5 :Evaluation accuracy on 300 images, 76.33
  19. ....................Loss tensor(1.8145, grad_fn=<DivBackward0>)
  20. Training: * Acc@1 58.833 Acc@5 82.333
  21. ..........Epoch 6 :Evaluation accuracy on 300 images, 74.33
  22. ....................Loss tensor(1.6930, grad_fn=<DivBackward0>)
  23. Training: * Acc@1 63.000 Acc@5 81.333
  24. ..........Epoch 7 :Evaluation accuracy on 300 images, 75.67

在这里，我们只对少数几个周期执行量化感知训练。 尽管如此，量化感知的训练在整个 imagenet 数据集上的准确率仍超过 71%，接近浮点精度 71.9%。

有关量化感知的训练的更多信息：

  • QAT 是训练后量化技术的超集，可以进行更多调试。 例如，我们可以分析模型的准确率是否受到权重或激活量化的限制。
  • 由于我们使用伪量化来对实际量化算术的数值建模，因此我们还可以在浮点中模拟量化模型的准确率。
  • 我们也可以轻松地模拟训练后量化。

来自量化的加速

最后，让我们确认一下我们上面提到的内容：量化模型实际上执行推理的速度更快吗？ 让我们测试一下：

  1. def run_benchmark(model_file, img_loader):
  2.     elapsed = 0
  3.     model = torch.jit.load(model_file)
  4.     model.eval()
  5.     num_batches = 5
  6.     # Run the scripted model on a few batches of images
  7.     for i, (images, target) in enumerate(img_loader):
  8.         if i < num_batches:
  9.             start = time.time()
  10.             output = model(images)
  11.             end = time.time()
  12.             elapsed = elapsed + (end-start)
  13.         else:
  14.             break
  15.     num_images = images.size()[0] * num_batches
  17.     print('Elapsed time: %3.0f ms' % (elapsed/num_images*1000))
  18.     return elapsed
  20. run_benchmark(saved_model_dir + scripted_float_model_file, data_loader_test)
  22. run_benchmark(saved_model_dir + scripted_quantized_model_file, data_loader_test)

出：

  1. Elapsed time:   7 ms
  2. Elapsed time:   4 ms

在 MacBook Pro 上本地运行此程序，常规模型的运行时间为 61 毫秒，而量化模型的运行时间仅为 20 毫秒，这说明了量化模型与浮点模型相比，典型的 2-4 倍加速。

总结

在本教程中，我们展示了两种量化方法-训练后静态量化和量化感知训练-描述它们在“幕后”进行的操作以及如何在 PyTorch 中使用它们。

谢谢阅读！ 与往常一样，我们欢迎您提供反馈，因此，如果有任何问题，请在这里创建一个 ISSUE。

脚本的总运行时间：（5 分钟 40.226 秒）

下载 Python 源码：static_quantization_tutorial.py

下载 Jupyter 笔记本：static_quantization_tutorial.ipynb

由 Sphinx 画廊生成的画廊