Implement Custom ReLU Autograd Function

easyFREECustom Autograd FunctionsPyTorch 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). Use the convention that the gradient is $0$ where the input is $\le 0$ (including exactly $0$ ) and passes the incoming gradient through where the input is $> 0$ .

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

Input: the harness builds a fixed 1-D vector $x = [-2, -1, 0, 1, 2]$ that tracks gradients, so the function takes no arguments.
Forward: ReLU is $f(x) = \max(0, x)$ applied element by element: $\max(0,-2)=0$ , $\max(0,-1)=0$ , $\max(0,0)=0$ , $\max(0,1)=1$ , $\max(0,2)=2$ , giving [0.0, 0.0, 0.0, 1.0, 2.0].
Upstream gradient: the loss is $L = \sum_i f(x_i)$ , so $\frac{\partial L}{\partial f(x_i)} = 1$ for every entry; the incoming gradient is all ones.
Backward: by the chain rule, $\frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial f(x_i)} \cdot f'(x_i)$ , where $f'(x) = 1$ for $x > 0$ and $0$ for $x \le 0$ . That yields $1 \cdot [0, 0, 0, 1, 1]$ = [0.0, 0.0, 0.0, 1.0, 1.0].
Notice: ReLU has no true derivative at $x = 0$ ; the convention here treats it as $0$ , so the third gradient is 0.0 even though that input sits on the kink. The derivative is $1$ exactly where the output is positive and $0$ where it is $0$ , so the gradient is a $0/1$ pattern scaled by the incoming gradient.

Constraints

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

Learn this in

PyTorch FundamentalsFoundations

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

Concept primer

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