📘

Implement Custom ReLU Autograd Function

EasyCustom Autograd Functions 📚 PyTorch Advanced

Problem Statement

Implement a custom ReLU activation function using torch.autograd.Function.

Background

PyTorch allows creating custom differentiable operations by subclassing torch.autograd.Function. You must implement:

forward(ctx, input): Compute the output and save tensors needed for backward
backward(ctx, grad_output): Compute gradients w.r.t. inputs

Your Task

The starter code defines a CustomReLU class and test harness. Implement the forward method (compute ReLU and save what's needed for backpropagation) and the backward method (compute the gradient of ReLU with respect to its input).

Output Format

The function returns a dictionary with "output" (ReLU values) and "grad" (input gradients).

Example:

Input:

None

Output:

{'output': [0.0, 0.0, 0.0, 1.0, 2.0], 'grad': [0.0, 0.0, 0.0, 1.0, 1.0]}

Reasoning:

We define a custom ReLU function CustomReLU that inherits from torch.autograd.Function, implementing the forward and backward methods. In forward, the input is saved and the output is computed as $max(0, x)$ using input.clamp(min=0).
The input tensor [-2.0, -1.0, 0.0, 1.0, 2.0] with requires_grad=True is created and passed through the custom ReLU function, resulting in the output [0.0, 0.0, 0.0, 1.0, 2.0].
The .sum().backward() method is called on the output, which computes the gradients of the input tensor. In the backward method, the gradient of the input is computed as $grad\_output$ where the input is greater than 0, and 0 otherwise, resulting in the gradient [0.0, 0.0, 0.0, 1.0, 1.0].
The output values and gradients are returned as a dictionary with keys "output" and "grad", resulting in the final output {'output': [0.0, 0.0, 0.0, 1.0, 2.0], 'grad': [0.0, 0.0, 0.0, 1.0, 1.0]}.

Constraints:

Must subclass torch.autograd.Function
Must use ctx.save_for_backward
Use @staticmethod for forward and backward

More from PyTorch Advanced

Implement SGD with MomentumHard Implement Custom SGD OptimizerHard Weight Decay for Specific LayersMedium Freeze Backbone, Train HeadMedium Per-Layer Learning RatesEasy

Background Knowledge

Introduction to Autograd

Autograd is a key component of PyTorch, allowing for automatic computation of gradients. It is essential for training neural networks, as it enables the computation of gradients of the loss with respect to the model's parameters. In PyTorch, autograd is implemented using a tape-based system, where all operations are recorded on a tape, and then the gradients are computed by backpropagating through this tape.

Custom Autograd Functions

PyTorch provides the ability to create custom autograd functions by subclassing torch.autograd.Function. This allows developers to define new, differentiable operations that can be used in their models. To create a custom autograd function, you must implement two methods: forward and backward. The forward method computes the output of the function and saves any tensors that are needed for the backward pass. The backward method computes the gradients of the output with respect to the inputs.

ReLU Activation Function

The ReLU (Rectified Linear Unit) activation function is a widely used activation function in neural networks. It is defined as $f(x) = \max(0, x)$ , which means that all negative values are set to 0, and all positive values are left unchanged. The derivative of the ReLU function is $f'(x) = \begin{cases} 0 & \text{if } x < 0 \\ 1 & \text{if } x > 0 \end{cases}$ . This derivative is used in the backward pass to compute the gradients of the loss with respect to the inputs.

Algorithm/Approach

The general approach to solving this problem involves:

Defining a custom autograd function that implements the ReLU activation function
Implementing the forward method to compute the output and save any necessary tensors
Implementing the backward method to compute the gradients of the output with respect to the inputs
Creating an input tensor and applying the custom ReLU function
Computing the gradients of the output with respect to the inputs using the backward method

Step-by-Step Strategy

To implement the solution, follow these steps:

Define a class that subclasses torch.autograd.Function and implements the ReLU activation function.
Implement the forward method to compute the output and save the input tensor using ctx.save_for_backward.
Implement the backward method to compute the gradients of the output with respect to the inputs.
Create an input tensor with requires_grad=True and apply the custom ReLU function.
Compute the gradients of the output with respect to the inputs using the backward method.
Return a dictionary with the output values and the gradients of the input.

Common Pitfalls

Some common pitfalls to watch out for when implementing custom autograd functions include:

Forgetting to save necessary tensors in the forward method
Incorrectly implementing the backward method
Not handling edge cases correctly (e.g., zero gradients for negative inputs)

Time & Space Complexity

The time complexity of the custom ReLU function is $O(n)$ , where $n$ is the number of elements in the input tensor. The space complexity is also $O(n)$ , as we need to save the input tensor in the forward method. The time and space complexity of the backward method are also $O(n)$ , as we need to compute the gradients of the output with respect to the inputs.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.