Pytorch custom activation functions?

Question

I\'m having issues with implementing custom activation functions in Pytorch, such as Swish. How should I go about implementing and using custom activation functions in Pytorch?<

Answer 1

There are four possibilities depending on what you are looking for. You will need to ask yourself two questions:

Q1) Will your activation function have learnable parameters?

If yes, you have no choice to create your activation function as an nn.Module class because you need to store those weights.

If no, you are free to simply create a normal function, or a class, depending on what is convenient for you.

Q2) Can your activation function be expressed as a combination of existing PyTorch functions?

If yes, you can simply write it as a combination of existing PyTorch function and won't need to create a backward function which defines the gradient.

If no you will need to write the gradient by hand.

Example 1: Swish function

The swish function f(x) = x * sigmoid(x) does not have any learned weights and can be written entirely with existing PyTorch functions, thus you can simply define it as a function:

def swish(x):
    return x * torch.sigmoid(x)

and then simply use it as you would have torch.relu or any other activation function.

Example 2: Swish with learned slope

In this case you have one learned parameter, the slope, thus you need to make a class of it.

class LearnedSwish(nn.Module):
    def __init__(self, slope = 1):
        super().__init__()
        self.slope = slope * torch.nn.Parameter(torch.ones(1))

    def forward(self, x):
        return self.slope * x * torch.sigmoid(x)

Example 3: with backward

If you have something for which you need to create your own gradient function, you can look at this example: Pytorch: define custom function