模型量化是为了减少模型的大小，以便在边缘设备上进行计算

首先建网络:

import torch
import torch.nn as nn
from torchsummary import summary
device = torch.device("cpu")
class SimpleNet(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(in_channels=12, out_channels=12, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2)
        self.conv3 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=3, stride=1, padding=1)
        self.conv4 = nn.Conv2d(in_channels=24, out_channels=24, kernel_size=3, stride=1, padding=1)
        self.fc = nn.Linear(in_features=16 * 16 * 24, out_features=num_classes)
    def forward(self, input):
        output = self.conv1(input)
        output = nn.ReLU()(output)
        output = self.conv2(output)
        output = nn.ReLU()(output)
        output = self.pool(output)
        output = self.conv3(output)
        output = nn.ReLU()(output)
        output = self.conv4(output)
        output = nn.ReLU()(output)
        output = output.view(-1, 16 * 16 * 24)
        output = self.fc(output)
        return output
model = SimpleNet().to(device=device)
print(model)

结果：

SimpleNet(
  (conv1): Conv2d(3, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv2): Conv2d(12, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv3): Conv2d(12, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv4): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (fc): Linear(in_features=6144, out_features=10, bias=True)
)

量化操作：

这里使用的动态量化：

quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)
print(quantized_model)

结果：

SimpleNet(
  (conv1): Conv2d(3, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv2): Conv2d(12, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv3): Conv2d(12, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv4): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (fc): DynamicQuantizedLinear(in_features=6144, out_features=10, dtype=torch.qint8, qscheme=torch.per_tensor_affine)
)

看下到底压缩了多少：

import os
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')

print_size_of_model(model)
print_size_of_model(quantized_model)

结果：

Size (MB): 0.287049
Size (MB): 0.103451

压缩的效果还是比较明显的。

量化只是对模型的权重数值进行了优化，上面的是权重值的对比，看起来还是不错的

summary(model.cuda(), input_size=(3, 512, 512))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 12, 512, 512]             336
            Conv2d-2         [-1, 12, 512, 512]           1,308
         MaxPool2d-3         [-1, 12, 256, 256]               0
            Conv2d-4         [-1, 24, 256, 256]           2,616
            Conv2d-5         [-1, 24, 256, 256]           5,208
            Linear-6                   [-1, 10]          61,450
================================================================
Total params: 70,918
Trainable params: 70,918
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 3.00
Forward/backward pass size (MB): 78.00
Params size (MB): 0.27
Estimated Total Size (MB): 81.27
----------------------------------------------------------------

从上面的数据可以看到，模型权重数值(Params size (MB): 0.27)只是占一部分 ，在这个模型中占用的比例很小

由于量化后的模型没法使用summary，所以没法直接观察

后端：x86、arm移动嵌入式平台；

对应参数：‘fbgemm’ 、 ‘qnnpack’

命令行：torch.quantization.get_default_qconfig('fbgemm')  

1.动态量化代码示例：

import torch

# define a floating point model
class M(torch.nn.Module):
    def __init__(self):
        super(M, self).__init__()
        self.fc = torch.nn.Linear(4, 4)

    def forward(self, x):
        x = self.fc(x)
        return x

# create a model instance
model_fp32 = M()
# create a quantized model instance
model_int8 = torch.quantization.quantize_dynamic(
    model_fp32,  # the original model
    {torch.nn.Linear},  # a set of layers to dynamically quantize
    dtype=torch.qint8)  # the target dtype for quantized weights

# run the model
input_fp32 = torch.randn(4, 4, 4, 4)
res = model_int8(input_fp32)

适用于 Linear、LSTM、RNN等层；

权重直接量化；bias和激活函数 在推理过程中动态量化；

2.静态量化示例：

import torch

# define a floating point model where some layers could be statically quantized
class M(torch.nn.Module):
    def __init__(self):
        super(M, self).__init__()
        # QuantStub converts tensors from floating point to quantized
        self.quant = torch.quantization.QuantStub()
        self.conv = torch.nn.Conv2d(1, 1, 1)
        self.relu = torch.nn.ReLU()
        # DeQuantStub converts tensors from quantized to floating point
        self.dequant = torch.quantization.DeQuantStub()

    def forward(self, x):
        # manually specify where tensors will be converted from floating
        # point to quantized in the quantized model
        x = self.quant(x)
        x = self.conv(x)
        x = self.relu(x)
        # manually specify where tensors will be converted from quantized
        # to floating point in the quantized model
        x = self.dequant(x)
        return x

# create a model instance
model_fp32 = M()

# model must be set to eval mode for static quantization logic to work
model_fp32.eval()

# attach a global qconfig, which contains information about what kind
# of observers to attach. Use 'fbgemm' for server inference and
# 'qnnpack' for mobile inference. Other quantization configurations such
# as selecting symmetric or assymetric quantization and MinMax or L2Norm
# calibration techniques can be specified here.
model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# Fuse the activations to preceding layers, where applicable.
# This needs to be done manually depending on the model architecture.
# Common fusions include `conv + relu` and `conv + batchnorm + relu`
model_fp32_fused = torch.quantization.fuse_modules(model_fp32, [['conv', 'relu']])

# Prepare the model for static quantization. This inserts observers in
# the model that will observe activation tensors during calibration.
model_fp32_prepared = torch.quantization.prepare(model_fp32_fused)

# calibrate the prepared model to determine quantization parameters for activations
# in a real world setting, the calibration would be done with a representative dataset
input_fp32 = torch.randn(4, 1, 4, 4)
model_fp32_prepared(input_fp32)

# Convert the observed model to a quantized model. This does several things:
# quantizes the weights, computes and stores the scale and bias value to be
# used with each activation tensor, and replaces key operators with quantized
# implementations.
model_int8 = torch.quantization.convert(model_fp32_prepared)

# run the model, relevant calculations will happen in int8
res = model_int8(input_fp32)

1.静态量化需要在模型起始和结束位置定义quant和dequant接口；

2.配置好后端

3.融合的层声明；一般是conv+relu；或者是conv+bn+relu；

3.准备量化

4.配置量化的推理数据集（一般对应于你的训练任务）

5.量化模型转换；此处转换为int8精度；

6.验证量化后模型；

订阅代码可关注：https://github.com/oyjGithub

pytorch现在已经支持模型量化，主要支持两种量化模式(qnnpack和fbgemm)：qnnpack是基于tensor的模式，这种相当于一种全局量化模式；另外fbgemm是基于通道的模式，相比qnnpack，fbgemm是以及基于通道的局部量化模式；需要说明的是这两种模式主要是针对权重而言的，激活函数的量化都是一样的。这里给出一个量化示例，pytorch目前只支持8bit的量化，通过自己实验发现对于浅层网络，分类精度损失不大，但是对于较深的网络量化就需要有针对性的优化，比如增大量化位宽，对部分权重或者activate分布不恰当的层加大这些层的位宽，或者模型训练过程中对权重或activate的大小进行约束，从而可以在量化的时候可以采用较低的位宽实现较高的精度保持。
实现过程依然是承接前一篇博客而来，只是添加了量化部分而已，代码实现如下：

import os
import torch
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from valid_simple_cnn5 import Net, QuantNet, get_dataloader, train, test, calibrate, print_size_of_model, time_model_evaluation


#---------------------------------------------------
#setting
input_shape = [1, 1, 28, 28]
BATCH_SIZE = 1
EPOCH = 10
GAMMA = 0.7
LR = 1.0
SAVE_MODEL = True
LOG_INTERVAL = 100
N_TEST_ITER = 3

QUANT_BACKEND = "fbgemm"  # fbgemm for x86 / qnnpack for arm
N_CALIBRATION = 10

SEED = 777
N_CPU = os.cpu_count()
DATA_PATH = "./"
PARAM_PATH = "mnist_qcnn05.pt"
SCRIPTED_PARAM_PATH = "mnist_jit_qcnn.pt"


'''USE_CUDA = True
GPU_ID = [3]
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join([str(i) for i in GPU_ID])
device = torch.device("cuda" if torch.cuda.is_available() and USE_CUDA else "cpu")
'''
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
torch.manual_seed(SEED)

TRAIN = False

#---------------------------------------------------
#training
model = Net().to(device)
print('model = ',model)
train_loader, test_loader = get_dataloader(DATA_PATH, BATCH_SIZE, N_CPU)

if os.path.exists(PARAM_PATH) and TRAIN == False:
    print('start load model...')
    model.load_state_dict(torch.load(PARAM_PATH,map_location=torch.device('cpu')))
    print('Loaded parameters from', PARAM_PATH)
else:
    print('start training...')
    optimizer = optim.Adadelta(model.parameters(), lr=LR)
    scheduler = StepLR(optimizer, step_size=1, gamma=GAMMA)
    for epoch in range(1, EPOCH + 1):
        train(model, device, train_loader, optimizer, epoch, LOG_INTERVAL)
        test(model, device, test_loader)
        scheduler.step()

        if SAVE_MODEL:
            torch.save(model.state_dict(), PARAM_PATH)


#---------------------------------------------------
#testing

print('testing acc...')
print('model = ',model)
test(model, "cpu", test_loader)
# print_size_of_model(model)
# time_model_evaluation(model, test_loader, N_TEST_ITER)

#--------------------------------------------------------------
#下面是基于pytorch 模型量化
#--------------------------------------------------------------
#层的合并
model = torch.quantization.fuse_modules(model,[["conv1","relu1"],
                                               ["conv2","relu2"],
                                               ["conv3","relu3"],
                                               ["conv4","relu4"],
                                               ["conv5","relu5"],
                                               ["conv6","relu6"],
                                               ["fc1","relu7"]])


#量化模式(qnnpack是基于tensor的量化,fbgemm是基于通道的量化)
QUANT_BACKEND = "qnnpack"
model.qconfig = torch.quantization.get_default_qconfig(QUANT_BACKEND)

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

#对模型进行统计获取activate分布相关的最大最小值
calibrate(model,train_loader,N_CALIBRATION)
print('calibration done')
#test(model, "cpu", test_loader)
print('model = ',model)


torch.save(model.state_dict(),'mnist_qcnn_fuse03.pt')

#转换为量化的模型
torch.quantization.convert(model,inplace=True)
print('model = ',model)
torch.save(model.state_dict(), 'mnist_qcnn_quant0.pt')

print('quant testing...')
test(model,'cpu',test_loader)

其他的代码可以通过链接下载：请戳
量化相关的东西后续会进一步更新和讲解，水平有限，不当之处请指教，希望对大家有帮助。