How to change all Relu to TLU in pretrained models?

[mobassir94]

i was trying to replace all the Relu activation functions with TLU of skresnet34 model as shown here : https://github.com/yukkyo/PyTorch-FilterResponseNormalizationLayer/blob/master/frn.py
and in filter response normalization layer paper.

i was trying this :

class TLU(nn.Module):
    def __init__(self, num_features = 32):
        """max(y, tau) = max(y - tau, 0) + tau = ReLU(y - tau) + tau"""
        super(TLU, self).__init__()
        self.num_features = num_features
        self.tau = nn.parameter.Parameter(
            torch.Tensor(1, num_features, 1, 1), requires_grad=True)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.zeros_(self.tau)

    def extra_repr(self):
        return 'num_features={num_features}'.format(**self.__dict__)

    def forward(self, x):
        return torch.max(x, self.tau)
    

def convert_relu_to_tlu(model):
    for child_name, child in model.named_children():
        if isinstance(child, nn.ReLU):
            setattr(model, child_name, TLU())
        else:
            convert_relu_to_tlu(child)

def getter(model, name):
    layer = model
    for attrib in name.split("."):
        layer = getattr(layer, attrib)
    return layer

def setter(model, name, layer):
    try:
        attrib, name = name.rsplit(".", 1)
        model = getter(model, attrib)
    except ValueError:
        pass
    setattr(model, name, layer)

for name, module in model.named_modules():
    if isinstance(module, nn.ReLU):
        relu = getter(model, name)
        print(relu)
        #tlu = TLU(num_features = relu.num_features)
        tlu = TLU()
        print("changing {} with {}".format(relu, tlu))

        setter(model, name, tlu)

it replaces all relu with TLU but i can't pass exact value of num_features in TLU while replacing relu with tlu,how to modify this line of code :

tlu = TLU(num_features = relu.num_features)

for changing all relu with tlu?

[rwightman]

@mobassir94 For most models here, with any standard activation you can just create_model(model_name, act_layer=TLU) ... however, your activation is non-standard in that it requires feature dim and that is not passed to any activation. Extracting that from previous layer might be possible if you walk in order, most models should have modules defined in order of use, but there can tensor slicing/cat ops between prev model and act that might change dim in some models... basically not easy

[mobassir94]

@rwightman any plan on releasing some pretrained models using frn+tlu? not everyone got fancy hardware and can't pass large batch size for reducing error rate,if you could release some pretrained models that uses frn+tlu then a lot of deep learning practitioner like me could do well,some of your pre trained models are very nice,i would like to use some of them on my hardware 3060ti(8gb vram)

at least if you could provide a step by step tutorial like how we can do frn and tlu on your pretrained weights then it would be very helpful.
if you prepare a step by step tutorial for a model like ViT large,then we can follow your instructions and try to do same thing on other architectures on our own!!
in segmentation models pytorch library qubvel released some weights that uses groupnorm and they performed really well on image segmentation task.
anyway thank you :)