Hi, how should I save the computation graph of a gradient vector computed from torch.autograd.grad(loss, model.parameters(), create_graph=True)?
The background is that I want to compute the Hessian-vector products of k vectors: H V, in which H is the Hessian of a neural network with n parameters, and V is a constant matrix with n rows and k columns.  To do that, I compute the gradient of the inner product between gradient of the network forward function g and V, with respect to the network parameters. An example that works for a tiny network is
import torch

# define the tiny "network"
class quadratic_fun(torch.nn.Module):
    def __init__(self):
        super(quadratic_fun, self).__init__()
        self.x = torch.nn.Parameter(torch.ones(5, requires_grad=True))
        self.y = torch.nn.Parameter(torch.ones(5, requires_grad=True))

    def forward(self):
        loss = torch.norm(self.x) ** 2 + torch.norm(self.y) ** 2
        return loss

# compute the flattened gradient with create_graph=True
model = quadratic_fun()
loss_quad = model.forward()
grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)
flat_grad = torch.cat([g.contiguous().view(-1) for g in grad_ft])

# generate the constant matrix V, and compute the matrix-gradient product
torch.manual_seed(0)
V = torch.randn((10, 3))
h = torch.matmul(flat_grad, V)

# compute the matrix-Jacobian product by iterating over the columns of the constant matrix
for i in range(3):
    hvp = torch.autograd.grad(h[i], model.parameters(), retain_graph=True)
    hvp_flat = torch.cat([g.contiguous().view(-1) for g in hvp])
    print(hvp_flat)

which gives
tensor([-2.2517, -0.8678, -0.6320, -2.5267,  0.2397, -0.2232, -0.9854,  0.2248,
        -0.2046,  0.1050])
tensor([-2.3047,  1.6974, -4.2304,  0.7000,  2.4753, -1.2272,  0.4968, -1.6821,
         1.5849,  1.0457])
tensor([-0.5012,  1.3840,  0.6445,  0.6163, -0.2869,  0.0632,  0.8794, -4.6321,
        -0.5793,  4.6044])

However, this is not feasible on CUDA when H is the Hessian of a large neural network: with retain_graph=True in the third from last line, the CUDA memory will quickly be filled up. While if I don’t retain the graph, the graph will be freed after one iteration of the for loop. In that case, I would need to compute the gradient again, which is time-consuming. Thus I wonder if I can save the not only the gradient value, but also its associated computation graph (both generated from grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)) to a file or buffer, and reload it in a later iteration of the for loop.
Some other posts I looked into but didn’t find an answer:


This post 2 suggests using JIT, but it is not clear to me how to use the API for the graph of a gradient vector.
A reply in this post 1 suggests to compute the matrix-Jacobian product with torch.autograd.functional.jacobian 1, but it looks like the API only works when the function to compute Jacobian is explicitly defined.)

Thanks!