如何為 PyTorch 2 Export 量化撰寫 Quantizer¶
建立於：2023 年 7 月 28 日 | 最後更新：2024 年 8 月 01 日 | 最後驗證：2024 年 11 月 05 日
作者: Leslie Fang, Weiwen Xia, Jiong Gong, Kimish Patel, Jerry Zhang
簡介¶
(prototype) PyTorch 2 Export Post Training 量化 介紹了 pytorch 2 export 量化的整體 API，與 fx graph mode 量化相比，API 的主要差異在於我們明確指出量化針對的是特定的後端。因此，要使用新的流程，後端需要實作一個 Quantizer 類別，該類別編碼：(1). 後端支援的量化運算符或模式是什麼 (2). 使用者如何表達他們希望如何量化浮點模型，例如，將整個模型量化為 int8 對稱量化，或僅量化線性層等。
請參閱 此處 了解新 API 和 Quantizer 的動機。
為 XNNPACK 定義的現有量化器物件位於 QNNPackQuantizer 中
註釋 API¶
Quantizer 使用註釋 API 來傳達不同運算符/模式的量化意圖。 註釋 API 主要包括 QuantizationSpec 和 QuantizationAnnotation。
QuantizationSpec 用於傳達張量將如何量化的意圖，例如 dtype、bitwidth、最小值、最大值、對稱與非對稱等。此外，QuantizationSpec 還允許量化器指定應如何觀察張量值，例如 MinMaxObserver 或 HistogramObserver 或某些自定義觀察器。
QuantizationAnnotation 由 QuantizationSpec 物件組成，用於註釋模式的輸入張量和輸出張量。 註釋輸入張量等同於註釋輸入邊緣，而註釋輸出張量等同於註釋節點。 QuantizationAnnotation 是一個具有多個欄位的 dataclass
input_qspec_map欄位屬於Dict類別，用於將每個輸入張量（作為輸入邊）對應到一個QuantizationSpec。output_qspec欄位表示用於註解輸出張量的QuantizationSpec；_annotated欄位指示此節點是否已由量化器註解過。
總之，註解 API 需要量化器註解圖的邊（輸入張量）或節點（輸出張量）。現在，我們將逐步介紹如何將註解 API 與不同類型的 QuantizationSpec 一起使用。
1. 註解常見運算符模式¶
為了使用量化的模式/運算符，例如 quantized add，後端開發人員將有意量化（由 QuantizationSpec 表示）模式的輸入和輸出。以下是一個範例流程（以 add 運算符為例），說明如何在具有註解 API 的量化工作流程中傳達此意圖。
步驟 1：識別 FX 圖中的原始浮點模式。有幾種方法可以識別此模式：量化器可以使用模式匹配器來匹配運算符模式；量化器可以從頭到尾遍歷節點，並比較節點的目標類型以匹配運算符模式。在此範例中，我們可以使用 get_source_partitions 來匹配此模式。原始浮點
add模式僅包含一個add節點。
add_partitions = get_source_partitions(gm.graph, [operator.add, torch.add])
add_partitions = list(itertools.chain(*add_partitions.values()))
for add_partition in add_partitions:
add_node = add_partition.output_nodes[0]
步驟 2：為模式的輸入和輸出定義
QuantizationSpec。QuantizationSpec定義了data type、qscheme以及關於使用者如何觀察或模擬量化張量的其他量化參數。
act_quantization_spec = QuantizationSpec(
dtype=torch.int8,
quant_min=-128,
quant_max=127,
qscheme=torch.per_tensor_affine,
is_dynamic=False,
observer_or_fake_quant_ctr=HistogramObserver.with_args(eps=2**-12),
)
input_act_qspec = act_quantization_spec
output_act_qspec = act_quantization_spec
步驟 3：使用
QuantizationAnnotation註解模式的輸入和輸出。在此範例中，我們將使用在上述步驟 2 中創建的QuantizationSpec，為add節點的兩個輸入和一個輸出創建QuantizationAnnotation物件。
input_qspec_map = {}
input_act0 = add_node.args[0]
input_qspec_map[input_act0] = input_act_qspec
input_act1 = add_node.args[1]
input_qspec_map[input_act1] = input_act_qspec
add_node.meta["quantization_annotation"] = QuantizationAnnotation(
input_qspec_map=input_qspec_map,
output_qspec=output_act_qspec,
_annotated=True,
)
在我們像這樣註解 add 節點之後，在後續的量化流程中，HistogramObserver 將在 prepare 階段插入到其兩個輸入節點和一個輸出節點中。並且 HistogramObserver 將在 convert 階段被替換為 quantize 節點和 dequantize 節點。
3. 使用固定的量化參數註解運算符¶
註解量化模型的另一個典型用例是用於量化參數事先已知的張量。例如，像 sigmoid 這樣的運算符，在輸入和輸出張量上具有預定義且固定的 scale/zero_point。FixedQParamsQuantizationSpec 專為此用例而設計。要使用 FixedQParamsQuantizationSpec，使用者需要顯式傳入 scale 和 zero_point 的參數。
步驟 1：識別 FX 圖中的原始浮點模式。我們可以使用
QuantizationSpec範例中介紹的相同方法來識別sigmoid模式。步驟 2：建立具有固定
scale、zero_point值的輸入的FixedQParamsQuantizationSpec物件。這些值將用於在轉換階段建立quantize節點和dequantize節點。步驟 3：註解輸入和輸出以使用此
FixedQParamsQuantizationSpec物件。
act_qspec = FixedQParamsQuantizationSpec(
dtype=torch.uint8,
quant_min=0,
quant_max=255,
qscheme=torch.per_tensor_affine,
scale=1.0 / 256.0,
zero_point=0,
)
sigmoid_node.meta["quantization_annotation"] = QuantizationAnnotation(
input_qspec_map={input_act: act_qspec},
output_qspec=act_qspec,
_annotated=True,
)
4. 使用導出的量化參數註解張量¶
另一個用例是定義張量的約束，其量化參數是從其他張量導出的。例如，如果我們要註解一個卷積節點，並將其 bias 輸入張量的 scale 定義為激活張量的 scale 和權重張量的 scale 的乘積。我們可以使用 DerivedQuantizationSpec 來註解此 conv 節點。
步驟 1：識別 FX 圖中的原始浮點模式。我們可以使用
QuantizationSpec範例中介紹的相同方法來識別convolution模式。步驟 2：定義
derive_qparams_fn函數，它接受ObserverOrFakeQuantize( ObserverBase 或 FakeQuantizeBase) 的列表作為輸入。從每個ObserverOrFakeQuantize物件，使用者可以獲取scale、zero point值。使用者可以根據從觀察器或偽量化實例計算出的量化參數，定義其關於如何導出新的scale、zero point值的啟發式方法。步驟 3：定義
DerivedQuantizationSpec物件，它接受以下輸入：EdgeOrNode物件的列表。對應於每個EdgeOrNode物件的觀察器將被傳遞到derive_qparams_fn函數中；derive_qparams_fn函數；幾個其他量化參數，例如dtype、qscheme。步驟 4：使用
QuantizationAnnotation註解此 conv 節點的輸入和輸出。
def derive_qparams_fn(obs_or_fqs: List[ObserverOrFakeQuantize]) -> Tuple[Tensor, Tensor]:
assert len(obs_or_fqs) == 2, \
"Expecting two obs/fqs, one for activation and one for weight, got: {}".format(len(obs_or_fq))
act_obs_or_fq = obs_or_fqs[0]
weight_obs_or_fq = obs_or_fqs[1]
act_scale, act_zp = act_obs_or_fq.calculate_qparams()
weight_scale, weight_zp = weight_obs_or_fq.calculate_qparams()
return torch.tensor([act_scale * weight_scale]).to(torch.float32), torch.tensor([0]).to(torch.int32)
bias_qspec = DerivedQuantizationSpec(
derived_from=[(input_act, node), (weight, node)],
derive_qparams_fn=derive_qparams_fn,
dtype=torch.int32,
quant_min=-2**31,
quant_max=2**31 - 1,
qscheme=torch.per_tensor_symmetric,
)
input_qspec_map = {input_act: act_quantization_spec, weight: weight_quantization_spec, bias: bias_qspec}
node.meta["quantization_annotation"] = QuantizationAnnotation(
input_qspec_map=input_qspec_map,
output_qspec=act_quantization_spec,
_annotated=True,
)
5. 帶有 Resnet18 的玩具示例¶
在使用 QuantizationAnnotation API 定義上述註解方法後，我們現在可以將它們放在一起以建構一個 BackendQuantizer 並使用 Torchvision Resnet18 執行一個 玩具示例。為了更好地理解最終示例，以下是示例中使用的類別和實用函數
QuantizationConfig 分別包含用於激活、權重和偏差的
QuantizationSpec。在註解模型時，可以使用 get_input_act_qspec、get_output_act_qspec、get_weight_qspec 和 get_bias_qspec 從
QuantizationConfig獲取特定模式的QuantizationSpec。
關於 PT2E 量化流程的 IR 的注意事項¶
IR 表示模型的中間表示形式，例如，torch IR (torch.nn 模組，torch.nn.functional 運算) 或 aten IR (torch.ops.aten.linear, …)。PT2E 量化流程正在使用 pre autograd aten IR ( torch.export API 的輸出)，以便我們支援訓練。如前所示，我們需要先匹配運算符或運算符模式，然後才能將註解附加到它們，那麼問題是我們如何匹配模式？
動機：直接匹配 aten IR 的問題¶
最直接的方法可能是直接匹配 aten IR。
範例
for n in gm.graph.nodes:
if n.op != "call_function" or n.target not in [
torch.ops.aten.relu.default,
torch.ops.aten.relu_.default,
]:
continue
relu_node = n
maybe_conv_node = n.args[0]
if (
not isinstance(maybe_conv_node, Node)
or maybe_conv_node.op != "call_function"
or maybe_conv_node.target
not in [
torch.ops.aten.conv1d.default,
torch.ops.aten.conv2d.default,
]
):
continue
# annotate conv and relu nodes
...
然而，使用此 IR 的一個問題是，如果模組或函數式操作的 PyTorch 實作發生變更，表示方式可能會發生變更。但這可能是出乎意料的，因為建模使用者通常假設，當 eager 模式模型程式碼沒有變更時，他們也應該在程式捕獲後獲得相同的模型表示。這個問題的一個具體影響是，如果 Quantizer 基於識別 aten IR 模式進行註釋，那麼它可能無法在 PyTorch 版本更新後識別該模式，並且相同的 eager 模式浮點數可能會保持未量化。
建議：使用 SubgraphMatcherWithNameNodeMap 進行模式匹配¶
因此，我們建議人們透過 SubgraphMatcherWithNameNodeMap（SubgraphMatcher 的改進版本，使其更容易查詢人們想要註釋的節點）來識別模式，方法是捕獲 torch IR 模式（使用與捕獲浮點模型相同的程式捕獲），而不是直接使用 aten IR 模式。
範例
def conv_relu_pattern(input, weight, bias):
conv = torch.nn.functional.conv2d(input, weight, bias)
output = torch.nn.functional.relu(conv)
# returns an additional dict that includes a map from name to node that we want to annotate
return relu, {"input": input, "weight": weight, "bias": bias, "output": output}
matcher = SubgraphMatcherWithNameNodeMap(conv_relu_pattern)
matches = matcher.match(model)
for match in matches:
# find input and output of the pattern
# annotate the nodes
name_node_map = match.name_node_map
input_node = name_node_map["input"]
weight_node = name_node_map["weight"]
bias_node = name_node_map["bias"]
output_node = name_node_map["relu"]
input_node.users[0].meta["quantization_annotation"] = ...
weight_node.users[0].meta["quantization_annotation"] = ...
bias_node.users[0].meta["quantization_annotation"] = ...
output_node.meta["quantization_annotation"] = ...
這樣一來，即使 nn 模組和函數式的實作發生變更，浮點模型的 aten IR 發生變更，Quantizer 仍然有效，但由於我們再次捕獲模式，而不是硬編碼模式的 aten IR，我們也會獲得更新的 aten IR，並且仍然能夠匹配模式。
需要注意的是，如果模式的輸入有多個使用者，除了檢查 aten op 目標之外，我們沒有好的方法來識別我們想要註釋哪個使用者節點。
另一個需要注意的是，我們需要確保有一個詳盡的範例列表（例如，2D、3D、4D 輸入，實際與符號輸入，training=True 與 training=False 等），以確保涵蓋從 torch IR 模式捕獲的不同可能的 aten IR 結果。
注意：我們可能會提供一些 (pattern, list of example_inputs) 或一些預先產生的 matcher 物件，以便人們將來可以直接使用它們。
結論¶
透過本教學，我們介紹了 PyTorch 2 中的新量化路徑。使用者可以了解如何使用 QuantizationAnnotation API 定義 BackendQuantizer，並將其整合到 PyTorch 2 Export Quantization 流程中。給出了 QuantizationSpec、SharedQuantizationSpec、FixedQParamsQuantizationSpec 和 DerivedQuantizationSpec 的範例，用於特定的註釋用例。您可以使用 XNNPACKQuantizer 作為範例，開始實作您自己的 Quantizer。之後，請按照 本教學 實際量化您的模型。
