L2 Gradient Penalty

mediumPROGradient ManipulationPyTorch Advanced

Problem Statement

Implement an L2 gradient penalty term, commonly used in Wasserstein GANs (WGAN-GP).

Background

The gradient penalty encourages the norm of gradients to be close to 1. Given an input, you compute gradients of the output w.r.t. the input and penalize when the norm deviates from 1.

Your Task

The starter code creates a model and input. Compute the gradient of the model's output with respect to the input (with graph creation enabled so you can backprop through the penalty). Then compute the L2 norm of that gradient and penalize its deviation from 1.

Output Format

Returns a dictionary with "output_value", "grad_norm", and "penalty".

Example

Input

None

Output

{'output_value': 1.2696, 'grad_norm': 0.9247, 'penalty': 0.0057}

Reasoning

We start by seeding the random number generator with torch.manual_seed(42) to ensure reproducibility, then create a linear model nn.Linear(3, 1) and an input tensor [[1.0, 2.0, 3.0]] with requires_grad=True.
The model output is computed as output = model(input), which applies a linear transformation to the input: $output = w \cdot input + b$ , where $w$ and $b$ are the model's weights and bias.
We then compute the gradients of the output with respect to the input using torch.autograd.grad() with create_graph=True, and calculate the L2 norm of these gradients: $grad\_norm = \sqrt{\sum_{i=1}^{3} grad_i^2}$ .
Finally, we calculate the gradient penalty as $(grad\_norm - 1)^2$ , which encourages the norm of gradients to be close to 1, and return the results in a dictionary with the output value, gradient norm, and penalty value, all rounded to 4 decimals.

Constraints

Use torch.autograd.grad with create_graph=True
Gradient penalty = (||grad||_2 - 1)^2
Input must have requires_grad=True

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 primerPREVIEW

Background Knowledge

The L2 Gradient Penalty is a term commonly used in Wasserstein GANs (WGAN-GP) to enforce a Lipschitz constraint on the discriminator. This constraint encourages the norm of gradients to be close to 1, which helps to stabilize the training process and improve the quality of generated samples. The Lipschitz constraint is a measure of how much the output of a function can change when the input changes. In the context of GANs, this constraint helps to prevent the discriminator from becoming too sensitive to small changes in the input.

The gradient penalty is computed by first calculating the gradients of the output with respect to the input using backpropagation. The L2 norm of these gradients is then computed, and the penalty term is calculated as the squared difference between the L2 norm and 1. This penalty term is added to the loss function of the discriminator, which encourages the discriminator to have a Lipschitz constant close to 1. The L2 norm is a measure of the magnitude of a vector, and it is defined as the square root of the sum of the squares of the elements of the vector.

In the context of PyTorch, the torch.autograd.grad() function can be used to compute the gradients of the output with respect to the input. The create_graph=True argument is used to create a new graph for the gradients, which allows for further backpropagation. The norm(2) method is used to compute the L2 norm of the gradients.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Define a model and an input tensor with requires_grad=True
Compute the output of the model using the input tensor
Compute the gradients of the output with respect to the input using torch.autograd.grad()
Compute the L2 norm of the gradients
Compute the gradient penalty term as the squared difference between the L2 norm and 1

This approach involves using PyTorch's autograd system to compute the gradients of the output with respect to the input, and then using these gradients to compute the gradient penalty term.

Step-by-Step Strategy

To implement the solution, follow these steps:

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results