fuse_fx¶
-
class
torch.quantization.quantize_fx.
fuse_fx
(model, fuse_custom_config_dict=None)[source]¶ Fuse modules like conv+bn, conv+bn+relu etc, model must be in eval mode. Fusion rules are defined in torch.quantization.fx.fusion_pattern.py
- Parameters
model (*) – a torch.nn.Module model
fuse_custom_config_dict (*) –
Dictionary for custom configurations for fuse_fx, e.g.:
fuse_custom_config_dict = { "additional_fuser_method_mapping": { (Module1, Module2): fuse_module1_module2 } # Attributes that are not used in forward function will # be removed when constructing GraphModule, this is a list of attributes # to preserve as an attribute of the GraphModule even when they are # not used in the code, these attributes will also persist through deepcopy "preserved_attributes": ["preserved_attr"], }
Example:
from torch.ao.quantization import fuse_fx m = Model().eval() m = fuse_fx(m)