fuse_fx

class torch.ao.quantization.quantize_fx.fuse_fx(model, fuse_custom_config=None, backend_config=None)

融合 conv+bn, conv+bn+relu 等模块，模型必须处于 eval 模式。融合规则在 torch.ao.quantization.fx.fusion_pattern.py 中定义。

参数

• model (*) – torch.nn.Module 模型

• fuse_custom_config (*) – fuse_fx 的自定义配置。有关更多详细信息，请参阅 FuseCustomConfig

返回类型

GraphModule

示例

from torch.ao.quantization import fuse_fx
m = Model().eval()
m = fuse_fx(m)