def run(self, num_iter, dummy=None):
if num_iter == 0:
return
lossvec = None
if self.debug:
lossvec = torch.zeros(num_iter + 1)
grad_mags = torch.zeros(num_iter + 1)
for i in range(num_iter):
self.x.requires_grad(True)
# Evaluate function at current estimate
self.f0 = self.problem(self.x)
# Compute loss
loss = self.problem.ip_output(self.f0, self.f0)
# Compute grad
grad = TensorList(torch.autograd.grad(loss, self.x))
# Update direction
if self.dir is None:
self.dir = grad
else:
self.dir = grad + self.momentum * self.dir
self.x.vdetach_()
self.x -= self.step_legnth * self.dir
if self.debug:
lossvec[i] = loss.item()
grad_mags[i] = sum(grad.view(-1) @ grad.view(-1)).sqrt().item()
if self.debug:
self.x.requires_grad_(True)
self.f0 = self.problem(self.x)
loss = self.problem.ip_output(self.f0, self.f0)
grad = TensorList(torch.autograd.grad(loss, self.x))
lossvec[-1] = self.problem.ip_output(self.f0, self.f0).item()
grad_mags[-1] = sum(
grad.view(-1) @ grad.view(-1)).cpu().sqrt().item()
self.losses = torch.cat((self.losses, lossvec))
self.gradient_mags = torch.cat((self.gradient_mags, grad_mags))
if self.plotting:
plot_graph(self.losses, self.fig_num[0], title='Loss')
plot_graph(self.gradient_mags,
self.fig_num[1],
title='Gradient magnitude')
self.x.vdetach_()
self.clear_temp()
def evaluate_CG_iteration(self, delta_x):
if self.analyze_convergence:
x = (self.x + delta_x).detach()
x.requires_grad_(True)
# compute loss and gradient
loss = self.problem(x)
grad = TensorList(torch.autograd.grad(loss, x))
# store in the vectors
self.losses = torch.cat(
(self.losses, loss.detach().cpu().view(-1)))
self.gradient_mags = torch.cat(
(self.gradient_mags,
sum(grad.view(-1)
@ grad.view(-1)).cpu().sqrt().detach().view(-1)))