For example I have simple net:
class Net(nn.Module):
def __init__(self):
super().__init__()
self.init = nn.Sequential(
nn.BatchNorm2d(3),
nn.Conv2d(3, 5, kernel_size=3),
nn.ReLU()
)
self.conv1 = nn.Conv2d(5, 10, kernel_size=1)
def forward(self, x):
x = self.init(x)
x = self.conv1(x)
return x
And I want to reinitialize activation layers like this:
net = Net()
for m in net.modules():
if type(m) == nn.ReLU:
m = nn.ELU()
But this code do nothing.
I can do it directly like this:
net.init[2] = nn.ELU()
but I want to get generic approach for this task. How I can do it?
Thanks in advance