I have a model that uses ReLU activation functions. I would like to replace every ReLU activation function of that model with another activation function.
I tried to iterate over the model using model.named_children()
and model.named_modules()
to find and replace the activation functions. However, they do not appear. I can only access the linear layers. The softmax function does also not appaer in the output. What do I do wrong? How can I access and replace relu() and softmax() of my model?
Here is my code:
import torch
import torch.nn as nn
import torch.functional as F
class MLP(nn.Module):
def __init__(self, num_in, num_hidden, num_out):
super().__init__()
self.linear1 = nn.Linear(num_in, num_hidden)
self.linear2 = nn.Linear(num_hidden, num_hidden)
self.linear3 = nn.Linear(num_hidden, num_out)
def forward(self, x):
x = torch.relu(self.linear1(x))
x = torch.relu(self.linear2(x))
x = F.softmax(self.linear3(x), dim=1)
return x
def main():
model = MLP(num_in=2, num_hidden=16, num_out=2)
print("named_children()")
for name, layer in model.named_children():
print(name, layer)
print("\nnamed_modules()")
for name, layer in model.named_modules():
print(name, layer)
if __name__ == "__main__":
main()
which produces the following output
named_children()
linear1 Linear(in_features=2, out_features=16, bias=True)
linear2 Linear(in_features=16, out_features=16, bias=True)
linear3 Linear(in_features=16, out_features=2, bias=True)
named_modules()
MLP(
(linear1): Linear(in_features=2, out_features=16, bias=True)
(linear2): Linear(in_features=16, out_features=16, bias=True)
(linear3): Linear(in_features=16, out_features=2, bias=True)
)
linear1 Linear(in_features=2, out_features=16, bias=True)
linear2 Linear(in_features=16, out_features=16, bias=True)
linear3 Linear(in_features=16, out_features=2, bias=True)