Motivation

深度學(xué)習(xí)模型為什么要量化
模型量化是深度學(xué)習(xí)Inference加速的關(guān)鍵技術(shù)之一着帽， 一般訓(xùn)練之后得到的模型采用float32 (fp32)格式存儲(chǔ)， 由于FP32 bit位數(shù)寬刁品， 而且浮點(diǎn)計(jì)算對(duì)硬件資源消耗高辱姨，造成內(nèi)存帶寬誊酌，模型吞吐率瓶頸。 量化(Quantization) 是解決FP32 的模型在內(nèi)存帶寬消耗椭员，推理速度的主要技術(shù)之一车海， 其采用定點(diǎn)(fixed point)或者整形數(shù)據(jù)(INT8)代替FP32類型， Hardware friendly隘击。

如何學(xué)習(xí)和掌握量化技術(shù)
模型量化涉及很多概念和算法侍芝，比如 對(duì)稱量化/非對(duì)稱量化， 線性量化/非線性量化 等等埋同， 了解這些基礎(chǔ)概念之后我們最好結(jié)合實(shí)踐這樣才能更好的理解和加深記憶竭贩。

本文以PyTorch提供的Quantization例子為例，分析簡單的模型量化過程莺禁。

Requirements

• Ubuntu 20.04
• PyTorch 1.11.0

量化的基本原理

量化本質(zhì)就是一個(gè)映射留量， 將一組浮點(diǎn)數(shù)映射為一組整形(例如INT8)的過程

• 量化: FP32 ---> INT8
• 反量化: INT8 ---> FP32

# 量化
xq = round(x_fp32 / scale + zero_point)

# 反量化
xqd = (xq - zero_point) * scale

量化參數(shù)scale, zero_point影響量化的精度

直接根據(jù)上面的公式驗(yàn)證一下:
上例子： 對(duì)一個(gè)FP32的數(shù)據(jù)進(jìn)行量化，觀察輸出結(jié)果

import torch
import numpy as np
import random
import matplotlib.pyplot as plt
import os

seed = 1234
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# 量化的基本原理
'''
# 量化
xq = round(x_fp32 / scale + zero_point)

# 反量化
xqd = (xq - zero_point) * scale

量化參數(shù)scale, zero_point影響量化的精度

'''

x = torch.rand(3, 4, dtype=torch.float32)

scale = 0.3
zero_point = 8
xq = torch.round((x / scale + zero_point))
xqd = (xq - zero_point) * scale

print('量化前的FP32 Tensor', x)
print('量化之后的INT8 Tensor', xq)
print('反量化之后的FP32 Tensor', xqd)

plt.stem(x.flatten().numpy())
# plt.subplot(1, 2, 2)
plt.stem(xqd.flatten().float().numpy())
plt.show()

由于量化參數(shù) scale, zero_point 是隨意選取的哟冬，因此反量化之后和原始的FP32數(shù)據(jù)存在一定的誤差：

image.png

采用PyTorch 提供的量化工具

PyTorch 提供了對(duì)Tensor 進(jìn)行量化的API函數(shù):

• 逐Tensor/Layer量化(Tensor/Channel wise Quantization): quantize_per_tensor(input: Tensor, scale: Tensor, zero_point: Tensor, dtype: _dtype) -> Tensor:
• 逐Channel量化 （Channel wise Quantization): quantize_per_channel(input: Tensor, scales: Tensor, zero_points: Tensor, axis: _int, dtype: _dtype) -> Tensor:

兩種量化的區(qū)別: Tensor-wise 和Channel-wise Quantization的主要區(qū)別是量化的粒度

• Tensor-wise: 粒度粗楼熄， 量化誤差相對(duì)大， 每個(gè)Tensor只有一個(gè)scale, zero_point 參數(shù)
• Chanel-wise: 粒度細(xì)浩峡， 量化誤差相對(duì)小可岂， Tensor的每個(gè)Channel都有獨(dú)自的scale, zp參數(shù)

Tensor-wise Quantization

# Tensor-wise 量化

def quantize_per_tensor():
    x = torch.rand(3, 4, dtype=torch.float32)
    # quantize， 量化參數(shù): scale + zp 保存在了QuantizeTensor
    xq = torch.quantize_per_tensor(x, scale=0.3, zero_point=8, dtype=torch.qint8)
    # 反量化
    xqd = torch.dequantize(xq)

    print('量化前的FP32 Tensor', x)
    print('量化之后的INT8 Tensor', xq)
    print('反量化之后的FP32 Tensor', xqd)

    # compare the x, xq
    # plt.subplot(1, 2, 1)
    plt.stem(x.flatten().numpy())
    # plt.subplot(1, 2, 2)
    plt.stem(xqd.flatten().float().numpy())
    plt.show()

quantize_per_tensor()

image.png

Channel-wise Quantization

# Channel-wise量化

def quantize_per_channel():
    x = torch.rand(3, 4, dtype=torch.float32)
    # quantize翰灾， 量化參數(shù): scale + zp 保存在了QuantizeTensor
    xq = torch.quantize_per_channel(x, scales=torch.tensor([0.3, 0.5, 0.4]), 
                                    zero_points=torch.tensor([8, 9, 10]),
                                    axis=0, dtype=torch.qint8)
    # 反量化
    xqd = torch.dequantize(xq)

    print('量化前的FP32 Tensor', x)
    print('量化之后的INT8 Tensor', xq)
    print('反量化之后的FP32 Tensor', xqd)

    # compare the x, xq
    # plt.subplot(1, 2, 1)
    plt.stem(x.flatten().numpy())
    # plt.subplot(1, 2, 2)
    plt.stem(xqd.flatten().float().numpy())
    plt.show()

quantize_per_channel()

image.png

對(duì)實(shí)際的CNN模型進(jìn)行量化

上面的例子是對(duì)單個(gè)Tensor進(jìn)行量化缕粹，方便掌握量化的思想稚茅； 下面采用實(shí)際的MobileNet-V2模型進(jìn)行量化.

量化前的準(zhǔn)備

ImageNet-2012 數(shù)據(jù)集準(zhǔn)備， 包含ImageNet-train, ImageNet-Val數(shù)據(jù)集平斩， 在實(shí)際的量化中需要進(jìn)行標(biāo)定(calibration) 以及精度測試亚享， 因此需要這兩個(gè)數(shù)據(jù)集

• ImageNet-train: 訓(xùn)練集，總共120w+ 圖像绘面， 量化的標(biāo)定過程會(huì)從imageNet-train中隨機(jī)選擇小部分?jǐn)?shù)據(jù)
• ImageNet-val: 驗(yàn)證集欺税， 包含5w張圖像， 用于測試量化前后模型精度(Top1 Acc) 的變化情況

MobileNet 量化實(shí)現(xiàn)

總體方法介紹：

對(duì)于一個(gè)CNN/RNN/Transformer等模型揭璃， 一般需要量化的對(duì)象包含2個(gè)部分:

• weight: 即模型訓(xùn)練之后的權(quán)重參數(shù)晚凿， 模型訓(xùn)練好之后這部分是固定的， 因此weight可以直接量化
• featureMap/Activation: 和輸入數(shù)據(jù)有關(guān)瘦馍， Activation指的是模型運(yùn)行期間的中間數(shù)據(jù)Intermediate data歼秽，如果需要對(duì)Activation進(jìn)行量化，則需要采用一部分代表性的輸入樣本產(chǎn)生中間數(shù)據(jù)情组，通過對(duì)中間數(shù)據(jù)進(jìn)行觀察量化燥筷，這個(gè)過程成為標(biāo)定Calibration

根據(jù)Activation是否需要離線量化，量化方法可分為2種:

1. 動(dòng)態(tài)量化 Dynamic Quantization: Activation的量化參數(shù)是在模型運(yùn)行時(shí)動(dòng)態(tài)確定的呻惕，這種方式得到的量化參數(shù)比較準(zhǔn)確荆责，缺點(diǎn)也很明顯: 運(yùn)行效率低滥比，加速效果不好
2. 靜態(tài)量化 Static Quantization: 即采用離線標(biāo)定的方法實(shí)現(xiàn)計(jì)算Activation的量化參數(shù)

下面的示例代碼采用Static Quantization:

1. 工程結(jié)構(gòu)：

   
   
   image.png

Utils: 一些PyTorch樣板代碼亚脆， 直接復(fù)用即可

image.png

utils.py

import os
import torch
import torchvision
import torchvision.transforms as transforms


# Helper functions
# 用于ImageNet 驗(yàn)證
class AverageMeter(object):
    def __init__(self, name, fmt=':f') -> None:
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
    
    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
    

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)


def accuracy(output, target, topk=(1,)):
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()

        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


def evaluate(model, criterion, data_loader, neval_batches, n_print=5):
    model.eval()
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    cnt = 0
    with torch.no_grad():
        for image, target in data_loader:
            output = model(image)
            # loss = criterion(output, target)
            cnt += 1
            acc1, acc5 = accuracy(output, target, topk=(1, 5))
            print('.', end = '')
            top1.update(acc1[0], n=image.size(0))
            top5.update(acc5[0], n=image.size(0))
            if cnt % n_print == 0:
                print(f'{top1}, {top5}')
            if cnt >= neval_batches:
                 return top1, top5

    return top1, top5


def load_model(model, model_weight_file):
    state_dict = torch.load(model_weight_file)
    model.load_state_dict(state_dict)
    model.to('cpu')
    return model


def print_size_of_model(model):
    torch.save(model.state_dict(), "temp.p")
    print('Size (MB):', os.path.getsize("temp.p")/1e6)
    os.remove('temp.p')


def prepare_data_loaders(data_path, train_batch_size, eval_batch_size):
    '''
    
    準(zhǔn)備ImageNet數(shù)據(jù)集， 包含train_set, val_set
    '''
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    
    dataset_train = torchvision.datasets.ImageNet(
        data_path, split='train',
        transform=transforms.Compose(
            [
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize
            ]))
    
    dataset_test = torchvision.datasets.ImageNet(
        data_path, split='val',
        transform=transforms.Compose(
            [
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize
            ]))

    train_sampler = torch.utils.data.RandomSampler(dataset_train)
    test_sampler = torch.utils.data.SequentialSampler(dataset_test)
    
    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, batch_size=train_batch_size,
        sampler=train_sampler
    )
    
    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, batch_size=eval_batch_size,
        num_workers=8, sampler=test_sampler
        )

    return data_loader_train, data_loader_test

2. MobileNet-V2模型

MV2模型基本上和torchvision中提供的模型寫法相同盲泛，只是個(gè)別地方的寫法需要修改濒持，為量化進(jìn)行適配

• Replace torch.add with floatfunctional， 為量化修改： self.skip_add = nn.quantized.FloatFunctional()
• QuantStub(), DeQuantStub(): 相當(dāng)于插樁寺滚， 告訴量化算法對(duì)模型的哪些區(qū)域進(jìn)行數(shù)據(jù)觀察以及量化

model.py

from statistics import mode
import numpy as np
import os
import matplotlib.pyplot as plt
import random
import torch
from torch import nn
import torchvision
from torch.quantization import QuantStub, DeQuantStub

# model = torchvision.models.mobilenet_v2(pretrained=True, progress=True)

# # Setup warnings
import warnings
warnings.filterwarnings(
    action='ignore',
    category=DeprecationWarning,
    module=r'.*'
)
warnings.filterwarnings(
    action='default',
    module=r'torch.quantization'
)

# https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html

seed = 1234
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)


# Define MobileNet

# 需要對(duì)MobileNetV2模型進(jìn)行小修改
# Replacing addition with nn.quantized.FloatFunctional
# Insert QuantStub and DeQuantStub at the beginning and end of the network.
# Replace ReLU6 with ReLU

def _make_divisible(v, divisor, min_value=None):
    """
    This function is taken from the original tf repo.
    It ensures that all layers have a channel number that is divisible by 8
    It can be seen here:
    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
    :param v:
    :param divisor:
    :param min_value:
    :return:
    """
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v


class ConvBNReLU(nn.Sequential):
    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
        padding = (kernel_size - 1) // 2
        super(ConvBNReLU, self).__init__(
            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
            nn.BatchNorm2d(out_planes, momentum=0.1),
            # Replace with ReLU
            nn.ReLU(inplace=False)
        )


class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = int(round(inp * expand_ratio))
        self.use_res_connect = self.stride == 1 and inp == oup

        layers = []
        if expand_ratio != 1:
            # pw
            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
        layers.extend([
            # dw
            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
            # pw-linear
            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
            nn.BatchNorm2d(oup, momentum=0.1),
        ])

        self.conv = nn.Sequential(*layers)
        # Replace torch.add with floatfunctional柑营， 為量化修改
        self.skip_add = nn.quantized.FloatFunctional()
    

    def forward(self, x):
        if self.use_res_connect:
            return self.skip_add.add(x, self.conv(x))   # 為量化修改
        else:
            return self.conv(x)


class MobileNetV2(nn.Module):
    def __init__(self, num_classes=1000, width_mult=1.0, inverted_residual_setting=None, round_nearest=8):
        """
        MobileNet V2 main class
        Args:
            num_classes (int): Number of classes
            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
            inverted_residual_setting: Network structure
            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
            Set to 1 to turn off rounding
        """
        super(MobileNetV2, self).__init__()
        block = InvertedResidual
        input_channel = 32
        last_channel = 1280

        if inverted_residual_setting is None:
            inverted_residual_setting = [
                # t, c, n, s
                [1, 16, 1, 1],
                [6, 24, 2, 2],
                [6, 32, 3, 2],
                [6, 64, 4, 2],
                [6, 96, 3, 1],
                [6, 160, 3, 2],
                [6, 320, 1, 1],
            ]

        # only check the first element, assuming user knows t,c,n,s are required
        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
            raise ValueError("inverted_residual_setting should be non-empty "
                             "or a 4-element list, got {}".format(inverted_residual_setting))

        # building first layer
        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
        features = [ConvBNReLU(3, input_channel, stride=2)]
        # building inverted residual blocks
        for t, c, n, s in inverted_residual_setting:
            output_channel = _make_divisible(c * width_mult, round_nearest)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
                input_channel = output_channel
        # building last several layers
        features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
        # make it nn.Sequential
        self.features = nn.Sequential(*features)
        self.quant = QuantStub()
        self.dequant = DeQuantStub()
        # building classifier
        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(self.last_channel, num_classes),
        )

        # weight initialization
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def forward(self, x):

        x = self.quant(x)

        x = self.features(x)
        x = x.mean([2, 3])
        x = self.classifier(x)
        x = self.dequant(x)
        return x
    
    
    # 模型Fuse優(yōu)化
    def fuse_model(self):
        for m in self.modules():
            if type(m) == ConvBNReLU:
                torch.quantization.fuse_modules(m, ['0', '1', '2'], inplace=True)
            if type(m) == InvertedResidual:
                for idx in range(len(m.conv)):
                    if type(m.conv[idx]) == nn.Conv2d:
                        torch.quantization.fuse_modules(m.conv, [str(idx), str(idx + 1)], inplace=True)


if __name__ == '__main__':
    net = MobileNetV2()
    net.eval()
    print(net)

    inputs = torch.rand(1, 3, 224, 224)
    outputs = net(inputs)
    print(outputs.size())

3. 量化測試

分別測試3種情況:

• 量化前模型的精度， baseline accuracy
• 采用layer-wise 進(jìn)行量化村视， 得到模型的量化后大小以及精度
• 采用Channel-wise進(jìn)行量化官套， 也得到量化后模型尺寸以及精度

baseline FP32模型精度測試
由于只是演示，而且筆者在CPU上跑模型蚁孔，因此只測試部分的樣本奶赔。
預(yù)訓(xùn)練模型從torchvision下載， 之后重命名為 mobilenet_pretrained_float.pth

代碼片段：

data_path = '/media/wei/Development/Datasets/imagenet'
saved_model_dir = 'data/'
float_model_file = 'mobilenet_pretrained_float.pth'
scripted_float_model_file = 'mobilenet_quantization_scripted.pth'
scripted_quantized_model_file = 'mobilenet_quantization_scripted_quantized.pth'

train_batch_size = 30
eval_batch_size = 512
num_eval_batches = 10

def evaluate_baseline():
    net = MobileNetV2()
    net.eval()
    float_model = load_model(net, saved_model_dir + float_model_file).to('cpu')

    print('\n Inverted Residual Block: Before fusion \n\n', float_model.features[1].conv)
    float_model.eval()

    # Fuses modules, Conv, BN, RelU算子融合
    float_model.fuse_model()

    # Note fusion of Conv+BN+Relu and Conv+Relu
    print('\n Inverted Residual Block: After fusion\n\n',float_model.features[1].conv)

    print("Size of baseline model")
    print_size_of_model(float_model)

    # ImageNet
    _, data_loader_test = prepare_data_loaders(data_path=data_path, train_batch_size=train_batch_size, eval_batch_size=eval_batch_size)

    # Start to get baseline Accuracy
    top1, top5 = evaluate(float_model, None, data_loader_test, neval_batches=num_eval_batches, n_print=1)
    print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))
    torch.jit.save(torch.jit.script(float_model), saved_model_dir + scripted_float_model_file)

Baseline模型的Top-1 Acc結(jié)果： Acc@Top1 = 78.57%

image.png

Layer-wise量化方法測試
需要設(shè)置量化方法: model.qconfig = torch.quantization.default_qconfig

查看一下default_qconfig的具體方法：

default_qconfig = QConfig(activation=default_observer,
weight=default_weight_observer)

• 對(duì)于Activation的量化方法: MinMax量化杠氢， 量化之后的值為0~127
  default_observer = MinMaxObserver.with_args(quant_min=0, quant_max=127)
• 對(duì)于Weight的量化方法： MinMax量化站刑， 量化為INT8, 對(duì)稱量化
  default_weight_observer = MinMaxObserver.with_args(
  dtype=torch.qint8, qscheme=torch.per_tensor_symmetric
  )

關(guān)于 MinMaxObserver: 統(tǒng)計(jì)一個(gè)Tensor中的最大和最小值， 因此量化方法為Tensor-wise或者稱為Layer-wise

class MinMaxObserver(_ObserverBase):
    r"""Observer module for computing the quantization parameters based on the
    running min and max values.

    This observer uses the tensor min/max statistics to compute the quantization
    parameters. The module records the running minimum and maximum of incoming
    tensors, and uses this statistic to compute the quantization parameters.

具體的測試代碼：

def per_layer_quantize():
    num_calibration_batches = 32
    
    net = MobileNetV2()
    net.eval()
    model = load_model(net, saved_model_dir + float_model_file).to('cpu')

    model.eval()
    # Fuses modules, Conv, BN, RelU算子融合
    model.fuse_model()

    # 設(shè)置量化配置鼻百，方法
    # MinMax observer, per-tensor quantize of weight
    model.qconfig = torch.quantization.default_qconfig
    print(model.qconfig)

    # model will be attached with qconfig
    torch.quantization.prepare(model, inplace=True)

    # Calibrate
    # 標(biāo)定應(yīng)該采用訓(xùn)練集的一部分
    data_loader_train, data_loader_test = prepare_data_loaders(data_path=data_path, eval_batch_size=eval_batch_size, train_batch_size=train_batch_size)

    # 會(huì)進(jìn)行標(biāo)定數(shù)據(jù)統(tǒng)計(jì)
    evaluate(model, criterion=None, data_loader=data_loader_train, neval_batches=num_calibration_batches)
    print('Calibration done...')

    # Convert to quantized model绞旅， Int8
    torch.quantization.convert(model, inplace=True)
    
    print('Quantize done')
    print('Size of quantized model')
    print_size_of_model(model)

    # test accuracy
    top1, top5 = evaluate(model, None, data_loader_test, neval_batches=num_eval_batches, n_print=1)
    print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))

運(yùn)行結(jié)果：
Acc@top1 = 65.37%摆尝， 比baseline低很多，精度丟失太嚴(yán)重Ｒ虮６楣！

image.png

Channel-wise 量化方法測試
由于Layer-wise量化之后模型的Top1 Acc下降太嚴(yán)重囤捻，因此需要更換量化方式臼朗。 Channel-wise是一種比layer-wise Quantization粒度更細(xì)的算法， 為Tensor的每個(gè)通道分別計(jì)算各自的量化參數(shù)蝎土，因此這個(gè)方法的精度預(yù)期比Layer-wise的高视哑。

Channel-wise量化的實(shí)現(xiàn)：
在PyTorch中為了支持不同的量化算法，Pytorch設(shè)置的不同的后端backend:

• 針對(duì)x86 CPU: fbgemm backend
• ARM CPU: qnnpack backend
  暫時(shí)未發(fā)現(xiàn)對(duì)GPU的支持誊涯，不過對(duì)于NVIDIA GPU, NVIDIA 開發(fā)的TensorRT應(yīng)該是支持的量化后端挡毅。

fbgemm backend:

def get_default_qconfig(backend='fbgemm'):
    """
    Returns the default PTQ qconfig for the specified backend.

    Args:
      * `backend`: a string representing the target backend. Currently supports `fbgemm`
        and `qnnpack`.

    Return:
        qconfig
    """

    if backend == 'fbgemm':
        qconfig = QConfig(activation=HistogramObserver.with_args(reduce_range=True),
                          weight=default_per_channel_weight_observer)
    elif backend == 'qnnpack':
        qconfig = QConfig(activation=HistogramObserver.with_args(reduce_range=False),
                          weight=default_weight_observer)
    else:
        qconfig = default_qconfig
    return qconfig

channel-wise的測試代碼：

def per_channel_quantize():
    num_calibration_batches = 32
    net = MobileNetV2()
    net.eval()
    model = load_model(net, saved_model_dir + float_model_file).to('cpu')
    model.eval()
    # Fuses modules, Conv, BN, RelU算子融合
    model.fuse_model()

    # 設(shè)置channel-wise 量化
    # qconfig = QConfig(activation=HistogramObserver.with_args(reduce_range=True),
    #               weight=default_per_channel_weight_observer)
    model.qconfig = torch.quantization.get_default_qconfig(backend='fbgemm')
    print('量化設(shè)置:', model.qconfig)

    torch.quantization.prepare(model, inplace=True)

     # Calibrate
    # 標(biāo)定應(yīng)該采用訓(xùn)練集的一部分
    data_loader_train, data_loader_test = prepare_data_loaders(data_path=data_path, eval_batch_size=eval_batch_size, train_batch_size=train_batch_size)

    # 會(huì)進(jìn)行標(biāo)定數(shù)據(jù)統(tǒng)計(jì)
    evaluate(model, criterion=None, data_loader=data_loader_train, neval_batches=num_calibration_batches)
    print('Calibration done...')

    # Convert to quantized model， Int8
    torch.quantization.convert(model, inplace=True)
    
    print('Quantize done')
    print('Size of quantized model')
    print_size_of_model(model)

    # test accuracy
    top1, top5 = evaluate(model, None, data_loader_test, neval_batches=num_eval_batches, n_print=1)
    print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))

量化結(jié)果：

• 模型大小下降的大約3~4倍： baseline FP32 13.99MB ---> 3.94MB
• 量化精度： baseline FP32=78.57%, layer-wise=65.37%, channel-wise=74.67%; 很顯然channel-wise的精度比layer-wise高很多暴构。

image.png

量化前后模型推理速度比較

CPU端測試跪呈， batch_size=4

• FP32 量化前的模型: 46ms
• INT8量化后的模型: 15ms, 推斷耗時(shí)僅為原來的1/3，加速效果非橙∮猓可觀

def speed_test():
    # Original FP32 model
    net = MobileNetV2()
    net.eval()
    model_fp32 = load_model(net, saved_model_dir + float_model_file).to('cpu')
    model_fp32.eval()
    # Fuses modules, Conv, BN, RelU算子融合
    model_fp32.fuse_model()

    # Int8 Quantized model
    model_quantized = per_channel_quantize(evaluate_acc=False)

    def inference_with_time_ms(BS=4, model=None):
        n = 10
        total_time = 0.0
        inputs = torch.rand(BS, 3, 224, 224)
        for i in range(n+2):
            t1 = time.time()
            model(inputs)
            torch.cuda.synchronize()
            t2 = time.time()
            if i > 1:
                total_time += (t2 - t1)
        return total_time / n * 1000

    print('Start speed test...')
    fp32_avg_time = inference_with_time_ms(BS=4, model=model_fp32)
    quantized_avg_time = inference_with_time_ms(BS=4, model=model_quantized)
    print(f'FP32 Avg time: {fp32_avg_time}ms')
    print(f'Quantized Avg time: {quantized_avg_time}ms')

Reference

• https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html