torch.ao.ns._numeric_suite

創建於: 2025年6月13日 | 最後更新於: 2025年6月13日

警告

This module is an early prototype and is subject to change.

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

Compare the weights of the float module with its corresponding quantized module. Return a dict with key corresponding to module names and each entry being a dictionary with two keys ‘float’ and ‘quantized’, containing the float and quantized weights. This dict can be used to compare and compute the quantization error of the weights of float and quantized models.

使用示例

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 of the float model

• quantized_dict (dict[str, Any]) – state dict of the quantized model

返回

dict with key corresponding to module names and each entry being a dictionary with two keys ‘float’ and ‘quantized’, containing the float and quantized weights

返回型別

weight_dict

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

Traverse the modules and save all logger stats into target dict. This is mainly used for quantization accuracy debug.

Type of loggers supported

ShadowLogger: used to log the outputs of the quantized module and its matching float shadow module, OutputLogger: used to log the outputs of the modules

引數

• mod (Module) – module we want to save all logger stats

• prefix (str) – prefix for the current module

返回

the dictionary used to save all logger stats

返回型別

target_dict

class torch.ao.ns._numeric_suite.Logger

Base class for stats logging

forward(x)

class torch.ao.ns._numeric_suite.ShadowLogger

Class used in Shadow module to record the outputs of the original and shadow modules.

forward(x, y)

class torch.ao.ns._numeric_suite.OutputLogger

Class used to log the outputs of the module

forward(x)

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

Shadow module attaches the float module to its matching quantized module as the shadow. Then it uses Logger module to process the outputs of both modules.

引數

• q_module – module quantized from float_module that we want to shadow

• float_module – float module used to shadow q_module

• logger_cls – type of logger used to process the outputs of q_module and float_module. ShadowLogger or custom loggers can be used.

forward(*x)

返回型別

張量

add(x, y)

返回型別

張量

add_scalar(x, y)

返回型別

張量

mul(x, y)

返回型別

張量

mul_scalar(x, y)

返回型別

張量

cat(x, dim=0)

返回型別

張量

add_relu(x, y)

返回型別

張量

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

Prepare the model by attaching the float module to its matching quantized module as the shadow if the float module type is in 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) – float module used to generate the q_module

• q_module (Module) – module quantized from float_module

• module_swap_list (set[type]) – list of float module types to attach the shadow

• logger_cls (Callable) – type of logger to be used in shadow module to process the outputs of quantized module and its 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'>)

Compare quantized module in a model with its floating point counterpart, feeding both of them the same input. Return 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 and its matching float shadow module. This dict can be used to compare and compute the module level quantization error.

This function first call prepare_model_with_stubs() to swap the quantized module that we want to compare with the Shadow module, which takes quantized module, corresponding float module and logger as input, and creates a forward path inside to make the float module to shadow quantized module sharing the same input. The logger can be customizable, default logger is ShadowLogger and it will save the outputs of the quantized module and float module that can be used to compute the module level quantization error.

使用示例

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) – float model used to generate the q_model

• q_model (Module) – model quantized from 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) – float module used to generate the q_module

• q_module (Module) – module quantized from 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)

透過將日誌記錄器附加到浮點模組和量化模組（如果它們在允許列表中）來準備模型。

引數

• float_module (Module) – float module used to generate the q_module

• q_module (Module) – module quantized from 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) – float model used to generate the q_model

• q_model (Module) – model quantized from float_model

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

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

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

返回

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

返回型別

act_compare_dict