How to Estimate the Gradient of a Function in One or More Dimensions in PyTorch?

Last Updated : 05 Jun, 2022

In this article, we are going to see how to estimate the gradient of a function in one or more dimensions in PyTorch.

torch.gradient() function

torch.gradient() method estimates the gradient of a function in one or more dimensions using the second-order accurate central differences method, and the function can be defined on a real or complex domain. For controllers and optimizers, gradient estimations are quite valuable. Gradient descent is a prominent optimization method that requires an estimate of the output derivatives with respect to each input at a given location. Let's have a look at the syntax of the given method first:

Syntax: torch.gradient(values)

Parameters:

values(Tensor): this parameter is represents the values of the function.

Example 1

In this example, we estimate the gradient of a function for a 1D tensor.

Python3

# Import required library
import torch

# define the tensor
tens = torch.tensor([-2., 1., -3., 4., 5.])
print(" Input tensor: ", tens)

# define a function
def fun(tens):
    return tens**2+5

# values of function
values = fun(tens)

# display values
print(" Function Values: ", values)

# estimate the gradients of fun
grad = torch.gradient(values)

# Display result
print(" Estimated Gradients of fun() - ", grad)

Output:

Example 2

In this example, we estimate the gradient of a function for a 2D tensor.

Python3

# Import required library
import torch

# define the tensor
tens = torch.tensor([[-1., 3., -5.],
                     [-4., 5.,  2.],
                     [-2., 3.,  4.], ])

print("\n Input tensor: \n", tens)

# define a function
def fun(tens):
    return tens**3

# values of function
values = fun(tens)

# display values
print("\n Function Values: \n", values)

# estimate the gradients of fun in dim=0
grad_dim_0 = torch.gradient(values, dim=0)
print("\n Estimated Gradients of fun() in dim=0 - \n", grad_dim_0)

# estimate the gradients of fun in dim=1
grad_dim_1 = torch.gradient(values, dim=1)
print("\n Estimated Gradients of fun() in dim=1 - \n", grad_dim_1)