Search code examples
pythonneural-networkdeep-learningpytorchactivation-function

How to change activation layer in Pytorch pretrained module?


How to change the activation layer of a Pytorch pretrained network? Here is my code :

print("All modules")
for child in net.children():
    if isinstance(child,nn.ReLU) or isinstance(child,nn.SELU):
        print(child)

print('Before changing activation')
for child in net.children():
    if isinstance(child,nn.ReLU) or isinstance(child,nn.SELU):
        print(child)
        child=nn.SELU()
        print(child)
print('after changing activation')
for child in net.children():
    if isinstance(child,nn.ReLU) or isinstance(child,nn.SELU):
        print(child)

Here is my output:

All modules
ReLU(inplace=True)
Before changing activation
ReLU(inplace=True)
SELU()
after changing activation
ReLU(inplace=True)

Solution

  • ._modules solves the problem for me.

    for name,child in net.named_children():
        if isinstance(child,nn.ReLU) or isinstance(child,nn.SELU):
            net._modules['relu'] = nn.SELU()