PyTorch数值工具套件教程

Created On: Jul 28, 2020 | Last Updated: Jan 16, 2024 | Last Verified: Not Verified

简介

量化很好，但只有在工作时才好；当它没有达到我们期望的准确性时，很难知道问题出在哪里。调试量化的准确性问题不易且耗时。

调试的重要一步是测量浮点模型及其对应的量化模型的统计数据，以了解它们的最大差异点。我们在PyTorch量化中构建了一套称为PyTorch Numeric Suite的数值工具，用于测量量化模块和浮点模块之间的统计数据，以支持量化调试工作。即使对具有良好准确性的量化模型，PyTorch Numeric Suite仍然可以作为剖析工具来更好地理解模型中的量化误差，为进一步优化提供指导。

PyTorch Numeric Suite目前支持通过静态量化和动态量化完成的模型，并具有统一的API。

在本教程中，我们将首先使用ResNet18作为示例展示如何使用PyTorch Numeric Suite测量静态量化模型和浮点模型之间的统计数据（即时模式下）。然后我们将使用基于LSTM的序列模型作为示例展示如何在动态量化模型中使用PyTorch Numeric Suite。

用于静态量化的数值工具套件

设置

我们先完成必要的导入操作：

import numpy as np
import torch
import torch.nn as nn
import torchvision
from torchvision import models, datasets
import torchvision.transforms as transforms
import os
import torch.quantization
import torch.quantization._numeric_suite as ns
from torch.quantization import (
    default_eval_fn,
    default_qconfig,
    quantize,
)

接着我们加载预训练的浮点ResNet18模型，并将其量化为qmodel。我们无法比较两个任意的模型，只能比较来自浮点模型的量化模型。

float_model = torchvision.models.quantization.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1, quantize=False)
float_model.to('cpu')
float_model.eval()
float_model.fuse_model()
float_model.qconfig = torch.quantization.default_qconfig
img_data = [(torch.rand(2, 3, 10, 10, dtype=torch.float), torch.randint(0, 1, (2,), dtype=torch.long)) for _ in range(2)]
qmodel = quantize(float_model, default_eval_fn, [img_data], inplace=False)

1. 比较浮点模型和量化模型的权重

我们通常首先想要比较的是量化模型和浮点模型的权重。我们可以调用PyTorch Numeric Suite中的``compare_weights()``获得一个字典``wt_compare_dict``，其中的键对应于模块名称，每个条目都是一个字典，拥有两个键'float'和'quantized'，分别包含浮点和量化权重。``compare_weights()``接收浮点和量化状态字典并返回一个字典，键对应于浮点权重，值是一个包含浮点和量化权重的字典。

wt_compare_dict = ns.compare_weights(float_model.state_dict(), qmodel.state_dict())

print('keys of wt_compare_dict:')
print(wt_compare_dict.keys())

print("\nkeys of wt_compare_dict entry for conv1's weight:")
print(wt_compare_dict['conv1.weight'].keys())
print(wt_compare_dict['conv1.weight']['float'].shape)
print(wt_compare_dict['conv1.weight']['quantized'].shape)

获得``wt_compare_dict``后，用户可以随意处理这个字典。下面作为示例，我们计算浮点模型和量化模型权重的量化误差。计算量化张量``y``的信号量化噪声比（SQNR）。SQNR反映了最大名义信号强度与量化中引入的量化误差之间的关系。更高的SQNR对应于更低的量化误差。

def compute_error(x, y):
    Ps = torch.norm(x)
    Pn = torch.norm(x-y)
    return 20*torch.log10(Ps/Pn)

for key in wt_compare_dict:
    print(key, compute_error(wt_compare_dict[key]['float'], wt_compare_dict[key]['quantized'].dequantize()))

作为另一个示例，``wt_compare_dict``也可用于绘制浮点模型和量化模型权重的直方图。

import matplotlib.pyplot as plt

f = wt_compare_dict['conv1.weight']['float'].flatten()
plt.hist(f, bins = 100)
plt.title("Floating point model weights of conv1")
plt.show()

q = wt_compare_dict['conv1.weight']['quantized'].flatten().dequantize()
plt.hist(q, bins = 100)
plt.title("Quantized model weights of conv1")
plt.show()

2. 在对应位置比较浮点模型和量化模型

第二个工具允许在相同输入情况下比较浮点和量化模型的权重和激活输出，并对应位置如图所示。红色箭头表示比较位置。

我们调用PyTorch Numeric Suite中的``compare_model_outputs()``以获取给定输入数据下在对应位置的浮点模型和量化模型的激活输出。此API返回一个字典，模块名称为键，每个条目本身是一个字典，拥有两个键'float'和'quantized'，分别包含激活输出。

data = img_data[0][0]

# Take in floating point and quantized model as well as input data, and returns a dict, with keys
# corresponding to the quantized module names and each entry being a dictionary with two keys 'float' and
# 'quantized', containing the activations of floating point and quantized model at matching locations.
act_compare_dict = ns.compare_model_outputs(float_model, qmodel, data)

print('keys of act_compare_dict:')
print(act_compare_dict.keys())

print("\nkeys of act_compare_dict entry for conv1's output:")
print(act_compare_dict['conv1.stats'].keys())
print(act_compare_dict['conv1.stats']['float'][0].shape)
print(act_compare_dict['conv1.stats']['quantized'][0].shape)

该字典可用于比较浮点模型和量化模型激活输出的量化误差，如下所示。

for key in act_compare_dict:
    print(key, compute_error(act_compare_dict[key]['float'][0], act_compare_dict[key]['quantized'][0].dequantize()))

如果我们想对多个输入数据进行比较，可以如下操作。通过在``white_list``中的情况下将记录器附加到浮点模块和量化模块来准备模型。默认记录器是``OutputLogger``，默认白名单是``DEFAULT_NUMERIC_SUITE_COMPARE_MODEL_OUTPUT_WHITE_LIST``。

ns.prepare_model_outputs(float_model, qmodel)

for data in img_data:
    float_model(data[0])
    qmodel(data[0])

# Find the matching activation between floating point and quantized modules, and return a dict with key
# corresponding to quantized module names and each entry being a dictionary with two keys 'float'
# and 'quantized', containing the matching floating point and quantized activations logged by the logger
act_compare_dict = ns.get_matching_activations(float_model, qmodel)

上述API中使用的默认记录器是``OutputLogger``，用于记录模块的输出。我们可以继承基类``Logger``并创建自己的记录器以执行不同的功能。例如，我们可以创建一个新的``MyOutputLogger``类，如下所示。

class MyOutputLogger(ns.Logger):
    r"""Customized logger class
    """

    def __init__(self):
        super(MyOutputLogger, self).__init__()

    def forward(self, x):
        # Custom functionalities
        # ...
        return x

然后可以将此记录器传递给上述API，例如：

data = img_data[0][0]
act_compare_dict = ns.compare_model_outputs(float_model, qmodel, data, logger_cls=MyOutputLogger)

或：

ns.prepare_model_outputs(float_model, qmodel, MyOutputLogger)
for data in img_data:
    float_model(data[0])
    qmodel(data[0])
act_compare_dict = ns.get_matching_activations(float_model, qmodel)

3. 将量化模型中的一个模块与其对应浮点版本进行比较，并使用相同输入数据

第三个工具允许将模型中量化模块与其对应的浮点模块进行比较，使用相同输入进行比较输出，如下图所示。

实际上，我们调用``prepare_model_with_stubs()``来将我们想要比较的量化模块替换为Shadow模块，如下所示：

Shadow模块以量化模块、浮点模块和记录器作为输入，并创建内部执行的前向路径，使浮点模块在共享相同输入张量的情况下模拟量化模块。

记录器可以自定义，默认记录器是``ShadowLogger``，它将保存量化模块和浮点模块的输出，可用于计算模块级别的量化误差。

注意在每次调用``compare_model_outputs()``和``compare_model_stub()``之前，我们需要有一个干净的浮点模型和量化模型。这是因为``compare_model_outputs()``和``compare_model_stub()``会内修改浮点模型和量化模型，并且如果一个紧接着另一个调用会导致意外结果。

float_model = torchvision.models.quantization.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1, quantize=False)
float_model.to('cpu')
float_model.eval()
float_model.fuse_model()
float_model.qconfig = torch.quantization.default_qconfig
img_data = [(torch.rand(2, 3, 10, 10, dtype=torch.float), torch.randint(0, 1, (2,), dtype=torch.long)) for _ in range(2)]
qmodel = quantize(float_model, default_eval_fn, [img_data], inplace=False)

在以下示例中，我们调用PyTorch Numeric Suite中的``compare_model_stub()``以比较``QuantizableBasicBlock``模块与其对应浮点版本。此API返回一个字典，其中的键对应于模块名称，每个条目是一个字典，具有两个键'float'和'quantized'，分别包含量化模块及其匹配浮点影子模块的输出张量。

data = img_data[0][0]
module_swap_list = [torchvision.models.quantization.resnet.QuantizableBasicBlock]

# Takes in floating point and quantized model as well as input data, and returns a dict with key
# corresponding to module names and each entry being a dictionary with two keys 'float' and
# 'quantized', containing the output tensors of quantized module and its matching floating point shadow module.
ob_dict = ns.compare_model_stub(float_model, qmodel, module_swap_list, data)

print('keys of ob_dict:')
print(ob_dict.keys())

print("\nkeys of ob_dict entry for layer1.0's output:")
print(ob_dict['layer1.0.stats'].keys())
print(ob_dict['layer1.0.stats']['float'][0].shape)
print(ob_dict['layer1.0.stats']['quantized'][0].shape)

然后可以使用该字典比较和计算模块级别的量化误差。

for key in ob_dict:
    print(key, compute_error(ob_dict[key]['float'][0], ob_dict[key]['quantized'][0].dequantize()))

如果我们想对多个输入数据进行比较，可以如下操作。

ns.prepare_model_with_stubs(float_model, qmodel, module_swap_list, ns.ShadowLogger)
for data in img_data:
    qmodel(data[0])
ob_dict = ns.get_logger_dict(qmodel)

上述API中使用的默认记录器是``ShadowLogger``，用于记录量化模块及其匹配浮点影子模块的输出。我们可以继承基类``Logger``并创建自己的记录器以执行不同的功能。例如，我们可以创建一个新的``MyShadowLogger``类，如下所示。

class MyShadowLogger(ns.Logger):
    r"""Customized logger class
    """

    def __init__(self):
        super(MyShadowLogger, self).__init__()

    def forward(self, x, y):
        # Custom functionalities
        # ...
        return x

然后可以将此记录器传递给上述API，例如：

data = img_data[0][0]
ob_dict = ns.compare_model_stub(float_model, qmodel, module_swap_list, data, logger_cls=MyShadowLogger)

或：

ns.prepare_model_with_stubs(float_model, qmodel, module_swap_list, MyShadowLogger)
for data in img_data:
    qmodel(data[0])
ob_dict = ns.get_logger_dict(qmodel)

用于动态量化的数值工具套件

Numeric Suite API设计成既适用于动态量化模型也适用于静态量化模型。我们将使用一个包含LSTM和线性模块的模型来展示Numeric Suite在动态量化模型上的使用。这个模型与动态量化LSTM语言模型教程中使用的模型一致 [1]。

设置

首先我们定义如下模型。注意在此模型中，只有``nn.LSTM``和``nn.Linear``模块会被动态量化，而``nn.Embedding``在量化后仍保持浮点模块。

class LSTMModel(nn.Module):
    """Container module with an encoder, a recurrent module, and a decoder."""

    def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):
        super(LSTMModel, self).__init__()
        self.encoder = nn.Embedding(ntoken, ninp)
        self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
        self.decoder = nn.Linear(nhid, ntoken)

        self.init_weights()

        self.nhid = nhid
        self.nlayers = nlayers

    def init_weights(self):
        initrange = 0.1
        self.encoder.weight.data.uniform_(-initrange, initrange)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-initrange, initrange)

    def forward(self, input, hidden):
        emb = self.encoder(input)
        output, hidden = self.rnn(emb, hidden)
        decoded = self.decoder(output)
        return decoded, hidden

    def init_hidden(self, bsz):
        weight = next(self.parameters())
        return (weight.new_zeros(self.nlayers, bsz, self.nhid),
                weight.new_zeros(self.nlayers, bsz, self.nhid))

随后我们创建``float_model``并将其量化为qmodel。

ntokens = 10

float_model = LSTMModel(
    ntoken = ntokens,
    ninp = 512,
    nhid = 256,
    nlayers = 5,
)

float_model.eval()

qmodel = torch.quantization.quantize_dynamic(
    float_model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)

1. 比较浮点模型和量化模型的权重

我们首先调用PyTorch Numeric Suite中的``compare_weights()``以获得一个字典``wt_compare_dict``，其中的键对应于模块名称，每个条目都是拥有两个键'float'和'quantized'的字典，分别包含浮点和量化权重。

wt_compare_dict = ns.compare_weights(float_model.state_dict(), qmodel.state_dict())

获得``wt_compare_dict``后，可以用于比较浮点模型和量化模型权重的量化误差，如下所示。

for key in wt_compare_dict:
    if wt_compare_dict[key]['quantized'].is_quantized:
        print(key, compute_error(wt_compare_dict[key]['float'], wt_compare_dict[key]['quantized'].dequantize()))
    else:
        print(key, compute_error(wt_compare_dict[key]['float'], wt_compare_dict[key]['quantized']))

上述``encoder.weight``条目中的Inf值是因为编码器模块未被量化，并且权重在浮点和量化模型中是相同的。

2. 在对应位置比较浮点模型和量化模型

然后我们调用PyTorch Numeric Suite中的``compare_model_outputs()``以获取给定输入数据下在对应位置的浮点模型和量化模型的激活输出。此API返回一个字典，模块名称为键，每个条目本身是一个字典，拥有两个键'float'和'quantized'，分别包含激活输出。注意此序列模型有两个输入，我们可以将两输入都传递给``compare_model_outputs()``和``compare_model_stub()``。

input_ = torch.randint(ntokens, (1, 1), dtype=torch.long)
hidden = float_model.init_hidden(1)

act_compare_dict = ns.compare_model_outputs(float_model, qmodel, input_, hidden)
print(act_compare_dict.keys())

该字典可用于比较浮点模型和量化模型激活输出的量化误差，如下所示。此模型中的LSTM模块有两个输出，在这个例子中我们计算第一个输出的误差。

for key in act_compare_dict:
    print(key, compute_error(act_compare_dict[key]['float'][0][0], act_compare_dict[key]['quantized'][0][0]))

3. 将量化模型中的一个模块与其对应浮点版本进行比较，并使用相同输入数据

接下来，我们调用 PyTorch 数字套件中的 compare_model_stub() 来比较 LSTM 和 Linear 模块与其浮点等效模块。该 API 返回一个字典，其键对应于模块名称，每个条目都是一个包含两个键的字典：’float’ 和 ‘quantized’，其中包含量化模块及其对应浮点影子模块的输出张量。

我们先重置模型。

float_model = LSTMModel(
    ntoken = ntokens,
    ninp = 512,
    nhid = 256,
    nlayers = 5,
)
float_model.eval()

qmodel = torch.quantization.quantize_dynamic(
    float_model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)

接下来，我们调用 PyTorch 数字套件中的 compare_model_stub() 来比较 LSTM 和 Linear 模块与其浮点等效模块。该 API 返回一个字典，其键对应于模块名称，每个条目都是一个包含两个键的字典：’float’ 和 ‘quantized’，其中包含量化模块及其对应浮点影子模块的输出张量。

module_swap_list = [nn.Linear, nn.LSTM]
ob_dict = ns.compare_model_stub(float_model, qmodel, module_swap_list, input_, hidden)
print(ob_dict.keys())

然后可以使用该字典比较和计算模块级别的量化误差。

for key in ob_dict:
    print(key, compute_error(ob_dict[key]['float'][0], ob_dict[key]['quantized'][0]))

40 dB 的 SQNR 是很高的，这表明浮点模型与量化模型之间有非常好的数值对齐。

总结

在本教程中，我们演示了如何使用 PyTorch 数字套件在动态模式下，通过统一的 API 测量和比较量化模型与浮点模型之间的统计数据，支持静态量化和动态量化。

感谢阅读！一如既往，我们欢迎任何反馈，如果您有任何问题，请在`此处 <https://github.com/pytorch/pytorch/issues>`_创建一个问题。

参考文献

[1] LSTM词语言模型的动态量化