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]) – 浮點模型的 state dict

• quantized_dict (dict[str, Any]) – 量化模型的 state dict

返回

字典，其鍵對應於模組名稱，每個條目是一個包含兩個鍵“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 量化而來且我們想要進行影子對比的模組

• float_module – 用於影子對比 q_module 的浮點模組

• logger_cls – 用於處理 q_module 和 float_module 輸出的 Logger 型別。可以使用 ShadowLogger 或自定義 Logger。

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)

如果浮點模組型別在 module_swap_list 中，則透過將其浮點模組作為影子附加到其匹配的量化模組來準備模型。

示例用法

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 的浮點模組

• q_module (Module) – 從 float_module 量化而來的模組

• module_swap_list (set[type]) – 要附加影子的浮點模組型別列表

• logger_cls (Callable) – 在影子模組中用於處理量化模組及其浮點影子模組輸出的日誌記錄器型別

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

比較模型中的量化模組與其對應的浮點模組，同時輸入相同的資料。返回一個字典，其鍵對應於模組名稱，每個條目是一個包含兩個鍵“float”和“quantized”的字典，分別包含量化模組及其匹配的浮點影子模組的輸出張量。此字典可用於比較和計算模組級別的量化誤差。

此函式首先呼叫 prepare_model_with_stubs() 將我們想要比較的量化模組替換為 Shadow 模組，Shadow 模組接受量化模組、對應的浮點模組和日誌記錄器作為輸入，並在內部建立一個前向路徑，使浮點模組成為量化模組的影子，共享相同的輸入。日誌記錄器可以自定義，預設日誌記錄器是 ShadowLogger，它將儲存量化模組和浮點模組的輸出，可用於計算模組級別的量化誤差。

示例用法

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 的浮點模型

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

• module_swap_list (set[type]) – 要附加影子模組的浮點模組型別列表。

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

• logger_cls – 在影子模組中用於處理量化模組及其浮點影子模組輸出的日誌記錄器型別

返回型別

dict[str, dict]

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

查詢浮點模組和量化模組之間的匹配啟用。

引數

• float_module (Module) – 用於生成 q_module 的浮點模組

• q_module (Module) – 從 float_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 的浮點模組

• q_module (Module) – 從 float_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 的浮點模型

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

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

• logger_cls – 要附加到 float_module 和 q_module 的日誌記錄器型別

• allow_list – 要附加日誌記錄器的模組型別列表

返回

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

返回型別

act_compare_dict