[FEATURE] Group Normalization for existing models

[sowmen]

Group Normalization seems to be effective for low batch training on less powerful h/w. Any plans to replace BN or add options for GN to existing models like EfficientNet? If not can you suggest a way I can implement GN to existing EfficientNet models? I was hoping to try the BiT (Big Transfer) technique.

[rwightman]

@sowmen I've added the ability to customize the norm layer and activation for a lot of the models in this collection, it shouldn't be too challenging. Actually adding such model defs and training them from scratch isn't currently in my plans.

Something (mostly), like below would work for EfficientNet

import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import GroupNorm
from timm.models.efficientnet import _gen_efficientnet
import timm.models
import functools


class GroupNorm(nn.GroupNorm):
    def __init__(self, num_channels, num_groups, eps=1e-5, affine=True):
        # NOTE num_channels is swapped to first arg for consistency in swapping norm layers with BN
        super().__init__(num_groups, num_channels, eps=eps, affine=affine)

    def forward(self, x):
        return F.group_norm(x, self.num_groups, self.weight, self.bias, self.eps)


@timm.models.register_model
def efficientnet_b0_gn(pretrained=False, **kwargs):
    """ EfficientNet-B0 """
    # NOTE for train, drop_rate should be 0.2, drop_path_rate should be 0.2
    gn = functools.partial(GroupNorm, num_groups=8)
    model = _gen_efficientnet(
        'efficientnet_b0', channel_multiplier=1.0, depth_multiplier=1.0,
        norm_layer=gn,  # change normalization layer
        pretrained=pretrained,
        external_default_cfg=dict(url=''),  # bit of a hack, make sure original pretrained weights don't get used
        **kwargs)
    return model

[mobassir94]

i would like to have this feature as well,,some models with frn layers with tlu could be great for people like us having low gpu vram,,replacing bn with frn+tlu for all different models are not straightforward task,,some of these pretrained models could help people like us a lot but no one working on it for pretraining models :(

[rwightman]

@mobassir94 it's a huge amount of time & work & hardware resources to define and train a series of new models w/ GN so it's unlikely to happen unless someone else takes it on. One of the likely reasons why this isn't common is that GN models are no faster, and actually use more GPU ram than a BN network. You can train w/ stability at smaller batch sizes, but they force you to use smaller batch sizes even sooner than BN based net.

[rwightman]

@hiyyg I'd probably just train a GN model myself as it's a data point I've wanted to explore with my current training recipes. In the meantime, the bitm resnetv2 models here are GN based with stdconv, the stdconv is much lower overhead than GN.

[hiyyg]

Sure. Then I will just use them without timm.

[rwightman]

Sure, I have some queued up to train now but it'll take a while, probably a few weeks. The bitm models are likely better than the detectron ones.

[hiyyg]

Not sure whether stdconv is really better, on my task, gn+ws is much worse than gn only.

[rwightman]

@hiyyg hmm, could also be the pre-activation, the v2 resnets outputs are pre-activation. Theoretically, the ws could be disabled and the values folded into the conv weights as a one time modification, then it'd just be a v2 + gn model for fine-tuning on other tasks, however that's a bit involved.