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)