Trouble with minimal hvp on pytorch model

user_650261 0
Viewed 256

While autograd's hvp tool seems to work very well for functions, once a model becomes involved, Hessian-vector products seem to go to 0. Some code.

First, I define the world's simplest model:

class SimpleMLP(nn.Module):
  def __init__(self, in_dim, out_dim):
      super().__init__()
      self.layers = nn.Sequential(
        nn.Linear(in_dim, out_dim),
      )
      
  def forward(self, x):
    '''Forward pass'''
    return self.layers(x)

Then, a loss function:

def objective(x):
  return torch.sum(0.25 * torch.sum(x)**4)

We instantiate it:

Arows = 2
Acols = 2

mlp = SimpleMLP(Arows, Acols)

Finally, I'm going to define a "forward" function (distinct from the model's forward function) that will serve as the the full model+loss that we want to analyze:

def forward(*params_list):
  for param_val, model_param in zip(params_list, mlp.parameters()):
    model_param.data = param_val
 
  x = torch.ones((Arows,))
  return objective(mlp(x))

This passes a ones vector into the single-layer "mlp," and passes it into our quadratic loss.

Now, I attempt to compute:

v = torch.ones((6,))
v_tensors = []
idx = 0
#this code "reshapes" the v vector as needed
for i, param in enumerate(mlp.parameters()):
  numel = param.numel()
  v_tensors.append(torch.reshape(torch.tensor(v[idx:idx+numel]), param.shape))
  idx += numel

And finally:

param_tensors = tuple(mlp.parameters())
reshaped_v = tuple(v_tensors)
soln =  torch.autograd.functional.hvp(forward, param_tensors, v=reshaped_v)

But, alas, the Hessian-Vector Product in soln is all 0's. What is happening?

autograd pytorch python
2 Answers

What's happening is that strict is False by default in the hvp() function and a tensor of 0's is returned as the Hessian Vector Product instead of an error (source).

If you try with strict=True, an error RuntimeError: The output of the user-provided function is independent of input 0. This is not allowed in strict mode. is returned instead. And when I looked at the full error, I suspect that this error comes from _check_requires_grad(jac, "jacobian", strict=strict) which indicates that the jacobian jac is None.

Update:

Following is a full working example:

import torch
from torch import nn

# your loss function
def objective(x):
    return torch.sum(0.25 * torch.sum(x)**4)

# Following are utilities to make nn.Module functional
# borrowed from the link I posted in comment
def del_attr(obj, names):
    if len(names) == 1:
        delattr(obj, names[0])
    else:
        del_attr(getattr(obj, names[0]), names[1:])

def set_attr(obj, names, val):
    if len(names) == 1:
        setattr(obj, names[0], val)
    else:
        set_attr(getattr(obj, names[0]), names[1:], val)

def make_functional(mod):
    orig_params = tuple(mod.parameters())
    # Remove all the parameters in the model
    names = []
    for name, p in list(mod.named_parameters()):
        del_attr(mod, name.split("."))
        names.append(name)
    return orig_params, names

def load_weights(mod, names, params):
    for name, p in zip(names, params):
        set_attr(mod, name.split("."), p)

# your forward function with update
def forward(*new_params):
    # this line replace your for loop
    load_weights(mlp, names, new_params)

    x = torch.ones((Arows,))
    out = mlp(x)
    loss = objective(out)
    return loss

# your simple MLP model
class SimpleMLP(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(in_dim, out_dim),
        )

    def forward(self, x):
        '''Forward pass'''
        return self.layers(x)


if __name__ == '__main__':
    # your model instantiation
    Arows = 2
    Acols = 2
    mlp = SimpleMLP(Arows, Acols)

    # your vector computation
    v = torch.ones((6,))
    v_tensors = []
    idx = 0
    #this code "reshapes" the v vector as needed
    for i, param in enumerate(mlp.parameters()):
        numel = param.numel()
        v_tensors.append(torch.reshape(torch.tensor(v[idx:idx+numel]), param.shape))
        idx += numel
    reshaped_v = tuple(v_tensors)

    #make model's parameters functional
    params, names = make_functional(mlp)
    params = tuple(p.detach().requires_grad_() for p in params)

    #compute hvp
    soln = torch.autograd.functional.vhp(forward, params, reshaped_v, strict=True)
    print(soln)
user_95609860

Did you try it with doubles instead of floats? I did some tests on my own that showed fairly large error when backproping with 32 bit float (on the order of 1e-5) compared to doubles.

user_4256460
Related