neural_compressor.adaptor.torch_utils.hawq_metric

Torch Utils for Hessian Aware Weighted Quantization.

Module Contents

Classes

Node_collector	Define Collector based on hook, which is used to record the intermediate result.
HessianTrace	HessianTrace Class.

Functions

compare_weights(→ Dict[str, Dict[str, torch.Tensor]])	Compare the weights of the float module with its corresponding quantized module.
hawq_top(fp32_model, q_model, dataloader, criterion, ...)	Enable hawq on an HessianTrace object and returns op list.

class neural_compressor.adaptor.torch_utils.hawq_metric.Node_collector(m)

Define Collector based on hook, which is used to record the intermediate result.

class neural_compressor.adaptor.torch_utils.hawq_metric.HessianTrace(model, dataloader, q_model, criterion=None)

HessianTrace Class.

Please refer to Yao, Zhewei, et al. “Pyhessian: Neural networks through the lens of the hessian.” 2020 IEEE international conference on big data (Big data). IEEE, 2020. Dong, Zhen, et al. “Hawq-v2: Hessian aware trace-weighted quantization of neural networks.” Advances in neural information processing systems 33 (2020): 18518-18529. https://github.com/openvinotoolkit/nncf/blob/develop/nncf/torch/quantization/hessian_trace.py

neural_compressor.adaptor.torch_utils.hawq_metric.compare_weights(float_dict: Dict[str, Any], quantized_dict: Dict[str, Any]) → Dict[str, Dict[str, torch.Tensor]]

Compare the weights of the float module with its corresponding quantized module.

Returns 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.

Example:

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

Parameters:

• float_dict – state dict of the float model.

• quantized_dict – state dict of the quantized model.

Returns:

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.

Return type:

weight_dict

neural_compressor.adaptor.torch_utils.hawq_metric.hawq_top(fp32_model, q_model, dataloader, criterion, enable_act)

Enable hawq on an HessianTrace object and returns op list.