Search code examples
pytorchneural-network

Is it possible to set different activation functions for different outputs at the final layer in the neural net?


I have a simple neural net - some linear layers with tanh between layers and after the end of the net. For example, I have input tensor with shape (100, 2) and I want the output to be with size (100, 5). But the values in the first column are in the range [0, 1], that is it is suitable to have sigmoid activation function at the end. The values in another columns are in the range [-1, 1], that is the "tanh" activation function could be used. But I don't understand how to set the sigmoid for the first output column and the "tanh" for another outputs? Is it possible? Or, I should apply abs() for the first column in the outputs and set the "tanh" after the final linear layer?

Now I have the following model:

nn.Sequential(nn.Linear(input_size, hidden_size),
nn.Tanh(),
nn.Linear(hidden_size, output_size),
nn.Tanh())

y = model(x)
y[:,0] = torch.abs(y[:,0])

But I want:

model = nn.Sequential(nn.Linear(input_size, hidden_size),
nn.Tanh(),
nn.Linear(hidden_size, output_size))

and to apply nn.Sigmoid() for the first output and nn.Tanh() for the other outputs:

y = model(x)
act_1 = nn.Sequential(nn.Sigmoid())
act_2 = nn.Sequential(nn.Tanh())

y[:,0] = act_1(y[:,0])
y[:,1:] = act_2(y[:,1:])

Solution

  • You shouldn't use Sequential for this. Just define your own class and implement your own custom logic, like so:

    import torch
    import torch.nn as nn
    
    class CustomModel(nn.Module):
        def __init__(self, input_size, hidden_size, output_size):
            super(CustomModel, self).__init__()
            self.fc1 = nn.Linear(input_size, hidden_size)
            self.fc2 = nn.Linear(hidden_size, output_size)
            self.tanh = nn.Tanh()
            self.sigmoid = nn.Sigmoid()
        
        def forward(self, x):
            x = self.tanh(self.fc1(x))
            
            x = self.fc2(x)
            
            # Sigmoid for the first column
            x[:, 0] = self.sigmoid(x[:, 0])
            # Tanh for the rest of the columns
            x[:, 1:] = self.tanh(x[:, 1:])
            
            return x
    
    input_size = 2
    hidden_size = 10
    output_size = 5
    model = CustomModel(input_size, hidden_size, output_size)
    
    x = torch.randn(100, 2)
    
    y = model(x)
    print(y)