Skip to content

fusion

fusion

Conv2d + BatchNorm2d fusion for quantization preparation.

fuse_conv_batchnorm 

fuse_conv_batchnorm(conv, batchnorm)

Create new Conv2d with BN folded into weights and bias.

Source code in src/versatil/post_training_compression/preparation/fusion.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def fuse_conv_batchnorm(
    conv: nn.Conv2d,
    batchnorm: nn.Module,
) -> nn.Conv2d:
    """Create new Conv2d with BN folded into weights and bias."""
    parameters = extract_batchnorm_parameters(batchnorm)
    if parameters is None:
        raise ValueError(
            f"Module {type(batchnorm).__name__} does not have "
            f"the required BatchNorm buffers"
        )
    running_mean, running_var, weight, bias, eps = parameters
    device = conv.weight.device
    scale = weight.to(device) / torch.sqrt(running_var.to(device) + eps)
    fused_conv = nn.Conv2d(
        in_channels=conv.in_channels,
        out_channels=conv.out_channels,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        dilation=conv.dilation,
        groups=conv.groups,
        bias=True,
        padding_mode=conv.padding_mode,
        device=device,
    )

    running_mean = running_mean.to(device)
    bias = bias.to(device)
    reshape_dims = (-1,) + (1,) * (conv.weight.dim() - 1)
    fused_conv.weight.data = conv.weight.data * scale.reshape(reshape_dims)
    if conv.bias is not None:
        fused_conv.bias.data = bias + (conv.bias.data - running_mean) * scale
    else:
        fused_conv.bias.data = bias - running_mean * scale
    return fused_conv

fuse_all_conv_batchnorm_pairs 

fuse_all_conv_batchnorm_pairs(model)

Find and fuse all consecutive Conv2d+BN children in all submodules.

Returns the number of pairs fused.

Source code in src/versatil/post_training_compression/preparation/fusion.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def fuse_all_conv_batchnorm_pairs(model: nn.Module) -> int:
    """Find and fuse all consecutive Conv2d+BN children in all submodules.

    Returns the number of pairs fused.
    """
    fusion_count = 0
    for parent in list(model.modules()):
        pairs = _find_conv_batchnorm_pairs(parent)
        for container, conv_name, batchnorm_name in pairs:
            conv = getattr(container, conv_name)
            batchnorm, activation = _unwrap_batchnorm_container(
                getattr(container, batchnorm_name)
            )
            fused = fuse_conv_batchnorm(conv=conv, batchnorm=batchnorm)
            setattr(container, conv_name, fused)
            if activation is not None:
                setattr(container, batchnorm_name, activation)
            else:
                setattr(container, batchnorm_name, nn.Identity())
            fusion_count += 1
    return fusion_count