torch.ao.ns._numeric_suite

警告

此模組是一個早期原型，可能會有所變動。

torch.ao.ns._numeric_suite.compare_weights(float_dict, quantized_dict)

比較浮點模組與其對應的量化模組的權重。返回一個字典，其鍵對應於模組名稱，每個條目都是一個字典，包含兩個鍵 ‘float’ 和 ‘quantized’，分別包含浮點和量化的權重。此字典可用於比較和計算浮點模型和量化模型的權重的量化誤差。

使用範例

wt_compare_dict = compare_weights(
    float_model.state_dict(), qmodel.state_dict())
for key in wt_compare_dict:
    print(
        key,
        compute_error(
            wt_compare_dict[key]['float'],
            wt_compare_dict[key]['quantized'].dequantize()
        )
    )

參數

• float_dict (Dict[str, Any]) – 浮點模型的狀態字典

• quantized_dict (Dict[str, Any]) – 量化模型的狀態字典

返回

字典，其鍵對應於模組名稱，每個條目都是一個字典，包含兩個鍵 ‘float’ 和 ‘quantized’，分別包含浮點和量化的權重

返回類型

weight_dict

torch.ao.ns._numeric_suite.get_logger_dict(mod, prefix='')

遍歷模組並將所有記錄器統計資料儲存到目標字典中。這主要用於量化準確性偵錯。

支援的記錄器類型

ShadowLogger：用於記錄量化模組及其匹配的浮點影子模組的輸出，OutputLogger：用於記錄模組的輸出

參數

• mod (Module) – 我們要儲存所有記錄器統計資料的模組

• prefix (str) – 目前模組的前綴

返回

用於儲存所有記錄器統計資料的字典

返回類型

target_dict

class torch.ao.ns._numeric_suite.Logger

統計資料記錄的基類

forward(x)

class torch.ao.ns._numeric_suite.ShadowLogger

在 Shadow 模組中使用的類別，用於記錄原始模組和影子模組的輸出。

forward(x, y)

class torch.ao.ns._numeric_suite.OutputLogger

用於記錄模組輸出的類別

forward(x)

class torch.ao.ns._numeric_suite.Shadow(q_module, float_module, logger_cls)

Shadow 模組將浮點模組附加到其匹配的量化模組作為影子。然後它使用 Logger 模組來處理兩個模組的輸出。

參數

• q_module – 從 float_module 量化的模組，我們想要 shadowing

• float_module – 用於 shadowing q_module 的浮點模組

• logger_cls – 用於處理 q_module 和 float_module 輸出的記錄器類型。可以使用 ShadowLogger 或自訂記錄器。

forward(*x)

返回類型

Tensor（張量）

add(x, y)

返回類型

Tensor（張量）

add_scalar(x, y)

返回類型

Tensor（張量）

mul(x, y)

返回類型

Tensor（張量）

mul_scalar(x, y)

返回類型

Tensor（張量）

cat(x, dim=0)

返回類型

Tensor（張量）

add_relu(x, y)

返回類型

Tensor（張量）

torch.ao.ns._numeric_suite.prepare_model_with_stubs(float_module, q_module, module_swap_list, logger_cls)

如果 float module 的類型在 module_swap_list 中，則將 float module 作為 shadow 連接到其匹配的量化 module，以此來準備模型。

使用範例

prepare_model_with_stubs(float_model, q_model, module_swap_list, Logger)
q_model(data)
ob_dict = get_logger_dict(q_model)

參數

• float_module (Module) – 用於生成 q_module 的 float module

• q_module (Module) – 從 float_module 量化的 module

• module_swap_list (Set[type]) – 要連接 shadow 的 float module 類型列表

• logger_cls (Callable) – 用於 shadow module 中的 logger 類型，以處理量化 module 及其 float shadow module 的輸出

torch.ao.ns._numeric_suite.compare_model_stub(float_model, q_model, module_swap_list, *data, logger_cls=<class 'torch.ao.ns._numeric_suite.ShadowLogger'>)

比較模型中量化的 module 與其對應的浮點 module，並為它們提供相同的輸入。返回一個 dict，其鍵對應於 module 名稱，每個條目都是一個包含兩個鍵“float”和“quantized”的字典，其中包含量化 module 及其匹配的 float shadow module 的輸出張量。該 dict 可用於比較和計算 module 級別的量化誤差。

此函數首先調用 prepare_model_with_stubs() 以將我們要比較的量化 module 與 Shadow module 交換，Shadow module 接收量化 module、相應的 float module 和 logger 作為輸入，並在內部創建一個前向路徑，使 float module 能夠對 shadow 量化 module 共享相同的輸入。 logger 可以是可自定義的，預設的 logger 是 ShadowLogger，它將保存量化 module 和 float module 的輸出，可用於計算 module 級別的量化誤差。

使用範例

module_swap_list = [torchvision.models.quantization.resnet.QuantizableBasicBlock]
ob_dict = compare_model_stub(float_model,qmodel,module_swap_list, data)
for key in ob_dict:
    print(key, compute_error(ob_dict[key]['float'], ob_dict[key]['quantized'].dequantize()))

參數

• float_model (Module) – 用於生成 q_model 的 float 模型

• q_model (Module) – 從 float_model 量化的模型

• module_swap_list (Set[type]) – 將附加 shadow module 的 float module 類型列表。

• data – 用於執行準備好的 q_model 的輸入資料

• logger_cls – 將在 shadow 模組中使用的記錄器類型，用於處理量化模組及其浮點 shadow 模組的輸出

返回類型

Dict[str, Dict]

torch.ao.ns._numeric_suite.get_matching_activations(float_module, q_module)

尋找浮點模組和量化模組之間的匹配激活。

參數

• float_module (Module) – 用於生成 q_module 的 float module

• q_module (Module) – 從 float_module 量化的 module

返回

字典，其鍵對應於量化模組名稱，每個條目都是一個字典，其中包含兩個鍵“float”和“quantized”，分別包含匹配的浮點和量化激活

返回類型

act_dict

torch.ao.ns._numeric_suite.prepare_model_outputs(float_module, q_module, logger_cls=<class 'torch.ao.ns._numeric_suite.OutputLogger'>, allow_list=None)

如果浮點模組和量化模組在 allow_list 中，則透過將記錄器附加到浮點模組和量化模組來準備模型。

參數

• float_module (Module) – 用於生成 q_module 的 float module

• q_module (Module) – 從 float_module 量化的 module

• logger_cls – 要附加到 float_module 和 q_module 的記錄器類型

• allow_list – 要附加記錄器的模組類型清單

torch.ao.ns._numeric_suite.compare_model_outputs(float_model, q_model, *data, logger_cls=<class 'torch.ao.ns._numeric_suite.OutputLogger'>, allow_list=None)

比較相同輸入下，浮點模型和量化模型在對應位置的輸出激活。 傳回一個字典，其鍵對應於量化模組名稱，每個條目都是一個字典，其中包含兩個鍵“float”和“quantized”，分別包含量化模型和浮點模型在匹配位置的激活。 此字典可用於比較和計算傳播量化誤差。

使用範例

act_compare_dict = compare_model_outputs(float_model, qmodel, data)
for key in act_compare_dict:
    print(
        key,
        compute_error(
            act_compare_dict[key]['float'],
            act_compare_dict[key]['quantized'].dequantize()
        )
    )

參數

• float_model (Module) – 用於生成 q_model 的 float 模型

• q_model (Module) – 從 float_model 量化的模型

• data – 用於執行準備好的 float_model 和 q_model 的輸入資料

• logger_cls – 要附加到 float_module 和 q_module 的記錄器類型

• allow_list – 要附加記錄器的模組類型清單

返回

字典，其鍵對應於量化模組名稱，每個條目都是一個字典，其中包含兩個鍵“float”和“quantized”，分別包含匹配的浮點和量化激活

返回類型

act_compare_dict