I am interested in computing the diagonal values of the Hessian of the loss function with respect to the parameters of an nn.Module to use for statistical inference (standard errors depend on the Hessian). The complication is that the model is estimated using mini-batches, because the full data doesn’t fit into memory, and existing solutions I’ve come across only give you the Hessian for the loss evaluated on a single batch of data.
Does anyone have ideas for how to compute the Hessian evaluated on the full dataset? For a single batch, a (not safe) method like the code below works. My current thinking is that I can just compute the Hessian for a loss function with `sum’ reduction on each batch separately and then sum across all batches, but I haven’t yet convinced myself this is valid so figured I’d see if someone else had thoughts
_input=torch.randn(32,3)
layer = nn.Linear(3,4)
criterion=nn.CrossEntropyLoss()

weight = layer.weight

def func(weight):
    del layer.weight
    layer.weight=weight
    return criterion(layer(_input), torch.zeros(len(_input),dtype=torch.long))

torch.autograd.functional.hessian(func,weight)

Related links:

Open issue with some discussion
Hackish solution for single batch 1
Slightly less hackish solution
Discussion on efficient computation of Jacobian