def f(flat_params_):
split_params = tensor_to_tuple(flat_params_, params)
load_weights(model, names, split_params)
out = model(x)
loss = model.loss(
out, y
) #calc_loss(out, y) + wd/2 * torch.sum(flat_params_*flat_params_)
return loss
def s_test_sample_exact(model, x_test, y_test, train_loader, gpu=-1):
grads = grad_z(x_test, y_test, model, gpu=gpu)
flat_grads = parameters_to_vector(grads)
def make_loss_f(model, params, names, x, y):
def f(flat_params_):
split_params = tensor_to_tuple(flat_params_, params)
load_weights(model, names, split_params)
out = model(x)
loss = model.loss(out, y)
return loss
return f
# Make model functional
params, names = make_functional(model)
# Make params regular Tensors instead of nn.Parameter
params = tuple(p.detach().requires_grad_() for p in params)
flat_params = parameters_to_vector(params)
h = torch.zeros([flat_params.shape[0], flat_params.shape[0]])
if gpu >= 0:
h = h.cuda()
# Compute real IHVP
for x_train, y_train in train_loader:
if gpu >= 0:
x_train, y_train = x_train.cuda(), y_train.cuda()
f = make_loss_f(model, params, names, x_train, y_train)
batch_h = hessian(f, flat_params, strict=True)
with torch.no_grad():
h += batch_h / float(len(train_loader))
h = (h + h.transpose(0,1))/2
with torch.no_grad():
load_weights(model, names, params, as_params=True)
inv_h = torch.inverse(h)
print("Inverse Hessian")
print(inv_h)
real_ihvp = inv_h @ flat_grads
return tensor_to_tuple(real_ihvp, params)
def f(flat_params_):
split_params = tensor_to_tuple(flat_params_, params)
load_weights(model, names, split_params)
out = model(x_train)
loss = calc_loss(out, y_train)
return loss
def f(flat_params_):
split_params = tensor_to_tuple(flat_params_, params)
load_weights(model, names, split_params)
out = model(x)
loss = model.loss(out, y)
return loss