how to get gradient of multiple outputs w.r.t. each input in a pytorch network?

Question

Assume that I have a simple neural network with 3 inputs x and 2 outputs y. the training samples would like to be :

x = torch.Tensor([[1,2,3],[4,5,6],[7,8,9]]) # each row is a sample
# [[a1,a2,a3],[b1,b2,b3],[c1,c2,c3]]

I want to get dy/dx for each one of the 2 outputs w.r.t each input in each sample. In other words, I want to get dy1/da1, dy1/da2 ...... dy1/dc3, dy2/da1......

I can hardly achieve this through torch.autograd.grad(). I constructed a code like this:

import torch
from torch import nn
from torch.nn.modules.activation import ReLU
import numpy as np

seed = 0

def seed_torch(seed=0):
    np.random.seed(seed)
    torch.manual_seed(seed) # set cpu random seed
    torch.cuda.manual_seed(seed) # set gpu random seed for single gpu
    torch.cuda.manual_seed_all(seed) # set gpus random seeds for multiple gpus
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

seed_torch()

x = torch.Tensor([[1,2,3],[4,5,6],[7,8,9]])
x.requires_grad = True

class model(nn.Module):
    def __init__(self):
        super(model, self).__init__()
        self.net = nn.Sequential()
        self.net.add_module('input_layer', nn.Linear(3, 10))
        self.net.add_module('activ_layer', nn.ReLU())
        self.net.add_module('output_layer', nn.Linear(10, 2))

    def forward(self, x):
        return self.net(x)

net = model()
row,col = net.net.input_layer.weight.data.shape
net.net.input_layer.weight.data = torch.Tensor(np.ones((row,col)))
num = net.net.input_layer.bias.data.shape
net.net.input_layer.bias.data = torch.Tensor(np.zeros(num))
row,col = net.net.output_layer.weight.data.shape
net.net.output_layer.weight.data = torch.Tensor(np.ones((row,col)))
num = net.net.output_layer.bias.data.shape
net.net.output_layer.bias.data = torch.Tensor(np.zeros(num))

y = net(x)
print(y)
# tensor([[ 60.,  60.],
#         [150., 150.],
#         [ 60.,  60.]], grad_fn=)
        
dydx = torch.autograd.grad(outputs=y, inputs=x,
    grad_outputs=torch.ones_like(y),
    retain_graph=True)[0]
print(dydx)
# tensor([[20., 20., 20.],
#         [20., 20., 20.],
#         [20., 20., 20.]])

dy1dx = [torch.autograd.grad(outputs=out, inputs=x,
    grad_outputs=torch.ones_like(out),
    retain_graph=True)[0][i] 
    for i, out in enumerate(y)]
print(dy1dx)
# [tensor([20., 20., 20.]), tensor([20., 20., 20.]), tensor([20., 20., 20.])]

print('finished')

but I think neither dydx nor dy1dx is what I want. I checked the docs and it says:

Computes and returns the sum of gradients of outputs with respect to the inputs.

It is suitable when the output is a scalar. But I may prefer to compute the gradients of vector or matrix outputs with respect to the inputs (maybe Jacobian matrix?), not the sum of them.

I also tried

dy2dx = torch.autograd.functional.jacobian(y, x, create_graph=True)
# or
dy2dx = torch.autograd.functional.jacobian(net(x), x, create_graph=True)
# meet: TypeError. 'Tensor' object is not callable
# or
dy2dx = torch.autograd.functional.jacobian(net(), x, create_graph=True)
# meet: TypeError. forward() missing 1 required positional argument: 'x'

but meet errors. Could anyone help me with this?

how to get gradient of multiple outputs w.r.t. each input in a pytorch network?

0 Answers