Skip to content

convert_layers

convert_layers

convert_layers 

convert_layers(model, layer_type_old, layer_type_new, convert_weights=False)

Recursively convert layers of a specific type in the model to a new type, in-place.

For GroupNorm conversion, num_groups is automatically computed as num_channels // 16 with fallback to ensure divisibility. Args: model: The model to modify in-place. layer_type_old: The type of layer to replace (e.g., nn.BatchNorm2d). layer_type_new: The type of layer to replace with (e.g., nn.GroupNorm). convert_weights: If True, copy weights and biases from old to new layers.

Returns:

Type Description
Module

The modified model with layers converted.
Source code in src/versatil/models/layers/convert_layers.py 
 5
 6
 7
 8
 9
10
11
12
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
49
50
51
52
53
54
55
56
def convert_layers(
    model: nn.Module,
    layer_type_old: type[nn.Module],
    layer_type_new: type[nn.Module],
    convert_weights: bool = False,
) -> nn.Module:
    """
    Recursively convert layers of a specific type in the model to a new type, in-place.

    For GroupNorm conversion, num_groups is automatically computed as num_channels // 16
    with fallback to ensure divisibility.
    Args:
        model: The model to modify in-place.
        layer_type_old: The type of layer to replace (e.g., nn.BatchNorm2d).
        layer_type_new: The type of layer to replace with (e.g., nn.GroupNorm).
        convert_weights: If True, copy weights and biases from old to new layers.

    Returns:
        The modified model with layers converted.
    """
    for name, module in reversed(model._modules.items()):
        if module is None:
            continue
        if len(list(module.children())) > 0:
            model._modules[name] = convert_layers(
                module, layer_type_old, layer_type_new, convert_weights
            )

        if isinstance(module, layer_type_old):
            layer_old = module
            num_channels = module.num_features

            # Compute num_groups automatically for GroupNorm
            if layer_type_new == nn.GroupNorm:
                computed_num_groups = _compute_num_groups(num_channels)
            else:
                computed_num_groups = num_channels

            layer_new = layer_type_new(
                num_groups=computed_num_groups,
                num_channels=num_channels,
                eps=module.eps,
                affine=module.affine,
            )

            if convert_weights and module.affine:
                layer_new.weight.data.copy_(layer_old.weight.data)
                layer_new.bias.data.copy_(layer_old.bias.data)

            model._modules[name] = layer_new

    return model

replace_batchnorm_with_groupnorm 

replace_batchnorm_with_groupnorm(model)

Replace all kinds of _BatchNorm layers (including SyncBatchNorm) in the model with GroupNorm layers.

Note: The number of groups for GroupNorm will be computed automatically as num_channels // 16.

Parameters:

Name Type Description Default
model Module

The model to modify

 required

Returns:

Type Description
Module

The modified model with BatchNorm replaced
Source code in src/versatil/models/layers/convert_layers.py 
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def replace_batchnorm_with_groupnorm(model: nn.Module) -> nn.Module:
    """Replace all kinds of `_BatchNorm` layers (including SyncBatchNorm) in the model with `GroupNorm` layers.

    Note: The number of groups for `GroupNorm` will be computed automatically as `num_channels // 16`.

    Args:
        model: The model to modify

    Returns:
        The modified model with `BatchNorm` replaced
    """
    return convert_layers(
        model,
        layer_type_old=_BatchNorm,
        layer_type_new=nn.GroupNorm,
        convert_weights=False,
    )