目錄
捷徑
prototype/backend_config_tutorial

(prototype) PyTorch BackendConfig 教學

建立於：2023 年 1 月 3 日 | 最後更新：2023 年 1 月 18 日 | 最後驗證：未驗證

作者: Andrew Or

BackendConfig API 使開發人員能夠將其後端與 PyTorch 量化整合。 目前僅在 FX 圖形模式量化中支援，但未來可能會將支援擴展到其他量化模式。 在本教學中，我們將示範如何使用此 API 客製化對特定後端的量化支援。 有關 BackendConfig 背後的動機和實作細節的更多資訊，請參閱此README。

假設我們是後端開發人員，並且希望將我們的後端與 PyTorch 的量化 API 整合。 我們的後端僅包含兩個運算子：量化線性和量化 conv-relu。 在本節中，我們將逐步說明如何透過 prepare_fx 和 convert_fx 使用自訂 BackendConfig 量化範例模型來實現此目的。

import torch
from torch.ao.quantization import (
    default_weight_observer,
    get_default_qconfig_mapping,
    MinMaxObserver,
    QConfig,
    QConfigMapping,
)
from torch.ao.quantization.backend_config import (
    BackendConfig,
    BackendPatternConfig,
    DTypeConfig,
    DTypeWithConstraints,
    ObservationType,
)
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx

1. 推導每個量化運算子的參考模式

對於量化線性，假設我們的後端期望參考模式 [dequant - fp32_linear - quant] 並將其降低為單個量化線性運算。 實現此目的的方法是首先在浮點線性運算前後插入 quant-dequant 運算，以便我們產生以下參考模型

quant1 - [dequant1 - fp32_linear - quant2] - dequant2

同樣，對於量化 conv-relu，我們希望產生以下參考模型，其中方括號中的參考模式將被降低為單個量化 conv-relu 運算

quant1 - [dequant1 - fp32_conv_relu - quant2] - dequant2

2. 使用後端約束設定 DTypeConfigs

在上面的參考模式中，DTypeConfig 中指定的輸入 dtype 將作為 dtype 引數傳遞給 quant1，而輸出 dtype 將作為 dtype 引數傳遞給 quant2。 如果輸出 dtype 是 fp32，就像動態量化的情況一樣，則不會插入輸出 quant-dequant 對。 此範例還顯示了如何指定對特定 dtype 的量化和縮放範圍的限制。

quint8_with_constraints = DTypeWithConstraints(
    dtype=torch.quint8,
    quant_min_lower_bound=0,
    quant_max_upper_bound=255,
    scale_min_lower_bound=2 ** -12,
)

# Specify the dtypes passed to the quantized ops in the reference model spec
weighted_int8_dtype_config = DTypeConfig(
    input_dtype=quint8_with_constraints,
    output_dtype=quint8_with_constraints,
    weight_dtype=torch.qint8,
    bias_dtype=torch.float)

3. 設定 conv-relu 的融合

請注意，原始使用者模型包含單獨的 conv 和 relu 運算，因此我們需要首先將 conv 和 relu 運算融合到單個 conv-relu 運算 (fp32_conv_relu) 中，然後以與量化線性運算類似的方式量化此運算。 我們可以透過定義一個接受 3 個引數的函數來設定融合，其中第一個引數是指是否用於 QAT，其餘引數指的是融合模式的各個項目。

def fuse_conv2d_relu(is_qat, conv, relu):
    """Return a fused ConvReLU2d from individual conv and relu modules."""
    return torch.ao.nn.intrinsic.ConvReLU2d(conv, relu)

4. 定義 BackendConfig

現在我們擁有了所有必要的元素，因此我們繼續定義我們的 BackendConfig。 在這裡，我們對線性運算的輸入和輸出使用不同的觀察器（將被重新命名），因此傳遞給兩個量化運算（quant1 和 quant2）的量化參數將有所不同。 對於像線性和 conv 這樣的加權運算，通常是這種情況。

對於 conv-relu 運算，觀察類型是相同的。 但是，我們需要兩個 BackendPatternConfigs 來支援此運算，一個用於融合，另一個用於量化。 對於 conv-relu 和線性，我們都使用上面定義的 DTypeConfig。

linear_config = BackendPatternConfig() \
    .set_pattern(torch.nn.Linear) \
    .set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) \
    .add_dtype_config(weighted_int8_dtype_config) \
    .set_root_module(torch.nn.Linear) \
    .set_qat_module(torch.nn.qat.Linear) \
    .set_reference_quantized_module(torch.ao.nn.quantized.reference.Linear)

# For fusing Conv2d + ReLU into ConvReLU2d
# No need to set observation type and dtype config here, since we are not
# inserting quant-dequant ops in this step yet
conv_relu_config = BackendPatternConfig() \
    .set_pattern((torch.nn.Conv2d, torch.nn.ReLU)) \
    .set_fused_module(torch.ao.nn.intrinsic.ConvReLU2d) \
    .set_fuser_method(fuse_conv2d_relu)

# For quantizing ConvReLU2d
fused_conv_relu_config = BackendPatternConfig() \
    .set_pattern(torch.ao.nn.intrinsic.ConvReLU2d) \
    .set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) \
    .add_dtype_config(weighted_int8_dtype_config) \
    .set_root_module(torch.nn.Conv2d) \
    .set_qat_module(torch.ao.nn.intrinsic.qat.ConvReLU2d) \
    .set_reference_quantized_module(torch.ao.nn.quantized.reference.Conv2d)

backend_config = BackendConfig("my_backend") \
    .set_backend_pattern_config(linear_config) \
    .set_backend_pattern_config(conv_relu_config) \
    .set_backend_pattern_config(fused_conv_relu_config)

5. 設定滿足後端約束的 QConfigMapping

為了使用上面定義的運算，使用者必須定義一個滿足 DTypeConfig 中指定的約束的 QConfig。 有關更多詳細資訊，請參閱 DTypeConfig 的文件。 然後，我們將對我們希望量化的模式中使用的所有模組使用此 QConfig。

# Note: Here we use a quant_max of 127, but this could be up to 255 (see `quint8_with_constraints`)
activation_observer = MinMaxObserver.with_args(quant_min=0, quant_max=127, eps=2 ** -12)
qconfig = QConfig(activation=activation_observer, weight=default_weight_observer)

# Note: All individual items of a fused pattern, e.g. Conv2d and ReLU in
# (Conv2d, ReLU), must have the same QConfig
qconfig_mapping = QConfigMapping() \
    .set_object_type(torch.nn.Linear, qconfig) \
    .set_object_type(torch.nn.Conv2d, qconfig) \
    .set_object_type(torch.nn.BatchNorm2d, qconfig) \
    .set_object_type(torch.nn.ReLU, qconfig)

6. 透過 prepare 和 convert 量化模型

最後，我們透過將我們定義的 BackendConfig 傳遞到 prepare 和 convert 中來量化模型。 這會產生一個量化的線性模組和一個融合的量化 conv-relu 模組。

class MyModel(torch.nn.Module):
    def __init__(self, use_bn: bool):
        super().__init__()
        self.linear = torch.nn.Linear(10, 3)
        self.conv = torch.nn.Conv2d(3, 3, 3)
        self.bn = torch.nn.BatchNorm2d(3)
        self.relu = torch.nn.ReLU()
        self.sigmoid = torch.nn.Sigmoid()
        self.use_bn = use_bn

    def forward(self, x):
        x = self.linear(x)
        x = self.conv(x)
        if self.use_bn:
            x = self.bn(x)
        x = self.relu(x)
        x = self.sigmoid(x)
        return x

example_inputs = (torch.rand(1, 3, 10, 10, dtype=torch.float),)
model = MyModel(use_bn=False)
prepared = prepare_fx(model, qconfig_mapping, example_inputs, backend_config=backend_config)
prepared(*example_inputs)  # calibrate
converted = convert_fx(prepared, backend_config=backend_config)

>>> print(converted)

GraphModule(
  (linear): QuantizedLinear(in_features=10, out_features=3, scale=0.012136868201196194, zero_point=67, qscheme=torch.per_tensor_affine)
  (conv): QuantizedConvReLU2d(3, 3, kernel_size=(3, 3), stride=(1, 1), scale=0.0029353597201406956, zero_point=0)
  (sigmoid): Sigmoid()
)

def forward(self, x):
    linear_input_scale_0 = self.linear_input_scale_0
    linear_input_zero_point_0 = self.linear_input_zero_point_0
    quantize_per_tensor = torch.quantize_per_tensor(x, linear_input_scale_0, linear_input_zero_point_0, torch.quint8);  x = linear_input_scale_0 = linear_input_zero_point_0 = None
    linear = self.linear(quantize_per_tensor);  quantize_per_tensor = None
    conv = self.conv(linear);  linear = None
    dequantize_2 = conv.dequantize();  conv = None
    sigmoid = self.sigmoid(dequantize_2);  dequantize_2 = None
    return sigmoid

(7. 實驗錯誤的 BackendConfig 設定)

作為一個實驗，這裡我們修改模型以使用 conv-bn-relu 而不是 conv-relu，但使用相同的 BackendConfig，它不知道如何量化 conv-bn-relu。 因此，只有 linear 被量化，而 conv-bn-relu 既沒有融合也沒有量化。

>>> print(converted)

GraphModule(
  (linear): QuantizedLinear(in_features=10, out_features=3, scale=0.015307803638279438, zero_point=95, qscheme=torch.per_tensor_affine)
  (conv): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
  (bn): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU()
  (sigmoid): Sigmoid()
)

def forward(self, x):
    linear_input_scale_0 = self.linear_input_scale_0
    linear_input_zero_point_0 = self.linear_input_zero_point_0
    quantize_per_tensor = torch.quantize_per_tensor(x, linear_input_scale_0, linear_input_zero_point_0, torch.quint8);  x = linear_input_scale_0 = linear_input_zero_point_0 = None
    linear = self.linear(quantize_per_tensor);  quantize_per_tensor = None
    dequantize_1 = linear.dequantize();  linear = None
    conv = self.conv(dequantize_1);  dequantize_1 = None
    bn = self.bn(conv);  conv = None
    relu = self.relu(bn);  bn = None
    sigmoid = self.sigmoid(relu);  relu = None
    return sigmoid

作為另一個實驗，這裡我們使用不滿足 backend 中指定的 dtype 限制的預設 QConfigMapping。 因此，沒有任何東西被量化，因為 QConfig 被直接忽略。

>>> print(converted)

GraphModule(
  (linear): Linear(in_features=10, out_features=3, bias=True)
  (conv): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
  (bn): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU()
  (sigmoid): Sigmoid()
)

def forward(self, x):
    linear = self.linear(x);  x = None
    conv = self.conv(linear);  linear = None
    bn = self.bn(conv);  conv = None
    relu = self.relu(bn);  bn = None
    sigmoid = self.sigmoid(relu);  relu = None
    return sigmoid

內建 BackendConfigs

PyTorch 量化支援一些內建的原生 BackendConfigs，位於 torch.ao.quantization.backend_config 命名空間下

還有其他正在開發中的 BackendConfigs（例如，用於 TensorRT 和 x86），但目前這些主要仍處於實驗階段。 如果使用者希望將新的、自定義的 backend 與 PyTorch 的量化 API 整合，他們可以使用與定義原生支援的 BackendConfigs 相同的 API 集合來定義自己的 BackendConfigs，如上例所示。

文件

取得 PyTorch 的全面開發者文件

檢視文件

教學課程

取得適用於初學者和進階開發人員的深入教學課程

檢視教學課程

資源

尋找開發資源並取得您問題的解答

檢視資源

為了分析流量並優化您的體驗，我們在此網站上提供 cookies。 點擊或導航即表示您同意允許我們使用 cookies。 作為本網站的目前維護者，Facebook 的 Cookie 政策適用。 了解更多資訊，包括有關可用控制項：Cookies 政策。