neural_compressor.adaptor.torch_utils.bf16_convert

Bf16 Convert for Torch Utils.

Module Contents

Classes

BF16ModuleWrapper

BF16Module Wrapper Class.

Functions

Convert(model, tune_cfg)

Convert to bf16 model.

bf16_symbolic_trace(model, fx_sub_module_list[, prefix])

Symbolic trace for bf16 models.
class neural_compressor.adaptor.torch_utils.bf16_convert.BF16ModuleWrapper(module)[source]

BF16Module Wrapper Class.

neural_compressor.adaptor.torch_utils.bf16_convert.Convert(model, tune_cfg)[source]

Convert to bf16 model.

Parameters:

  • model (object) – the input model.

  • tune_cfg (dict) – dictionary of quantization configuration.

Returns:

model with mixed precision.

Return type:

mixed_precision_model (object)

neural_compressor.adaptor.torch_utils.bf16_convert.bf16_symbolic_trace(model, fx_sub_module_list, prefix='')[source]

Symbolic trace for bf16 models.

Parameters:

  • model (object) – the input model.

  • fx_sub_module_list (list) – _description_

  • prefix (str) – prefix of op name.

Returns:

model (object)