fuse_fx

class torch.quantization.quantize_fx.fuse_fx(model, fuse_custom_config_dict=None)

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)