Shortcuts

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)

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources