How do you implement custom activation functions in PyTorch? Provide an example.?

In PyTorch, you can implement custom activation functions as a subclass of the ‘torch.autograd.Function‘ class, which allows the function to be used in a computational graph and have gradients computed for it during backpropagation.

Here’s an example implementation of a simple custom activation function called ‘my_activation‘ that applies the sigmoid function followed by a clamp operation to limit the output to the range [0.1, 0.9]:

import torch

class MyActivationFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        sigmoid_output = torch.sigmoid(input)
        clamped_output = torch.clamp(sigmoid_output, 0.1, 0.9)
        ctx.save_for_backward(input)
        return clamped_output

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        sigmoid_output = torch.sigmoid(input)
        grad_sigmoid = (sigmoid_output * (1 - sigmoid_output)) * grad_output
        grad_clamp = torch.zeros_like(grad_sigmoid)
        grad_clamp[(sigmoid_output >= 0.1) & (sigmoid_output <= 0.9)] = 1
        return grad_sigmoid * grad_clamp

my_activation = MyActivationFunction.apply

In the ‘forward‘ method, we compute the output of the activation function using the ‘sigmoid‘ and ‘clamp‘ operations, and save the input for use during backpropagation. In the ‘backward‘ method, we compute the gradients of the output with respect to the input using the chain rule, and also multiply by the gradient of the ‘clamp‘ operation which is 1 if the output is within the desired range and 0 otherwise.

To use this custom activation function in a neural network, you can just call it like any other PyTorch function:

import torch.nn as nn

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = my_activation(x)
        x = self.fc2(x)
        return x

In this example, we use our custom activation function after the first fully connected layer, and then connect it to the output layer. Whenever the network is backpropagated, gradients will be computed and propagated through the custom activation function.