Skip to content

bug: Batchnorms not replaced by Groupnorm when using share_rgb_model #151

@tlpss

Description

@tlpss

When using share_rgb_model to share the weights across different image inputs, the Batchnorms are not replaced by groupnorms in the MultiImageObsEncoder .

This results in very unstable training and reduced performance, as mentioned in the paper.

Can be fixed as follows:

        # handle sharing vision backbone
        if share_rgb_model:
            assert isinstance(rgb_model, nn.Module)
            key_model_map['rgb'] = rgb_model
            if use_group_norm:
                key_model_map['rgb'] = replace_submodules(
                    root_module=key_model_map['rgb'],
                    predicate=lambda x: isinstance(x, nn.BatchNorm2d),
                    func=lambda x: nn.GroupNorm(
                        num_groups=x.num_features//16, 
                        num_channels=x.num_features)
                )

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions