pytorch-L-BFGS-example

pytorch-lbfgs-example.py

import torch
import torch.optim as optim
import matplotlib.pyplot as plt


# 2d Rosenbrock function
def f(x):
    return (1 - x[0])**2 + 100 * (x[1] - x[0]**2)**2


# Gradient descent
x_gd = 10*torch.ones(2, 1)
x_gd.requires_grad = True
gd = optim.SGD([x_gd], lr=1e-5)
history_gd = []
for i in range(100):
    gd.zero_grad()
    objective = f(x_gd)
    objective.backward()
    gd.step()
    history_gd.append(objective.item())


# L-BFGS
def closure():
    lbfgs.zero_grad()
    objective = f(x_lbfgs)
    objective.backward()
    return objective

x_lbfgs = 10*torch.ones(2, 1)
x_lbfgs.requires_grad = True

lbfgs = optim.LBFGS([x_lbfgs],
                    history_size=10, 
                    max_iter=4, 
                    line_search_fn="strong_wolfe")
                    
history_lbfgs = []
for i in range(100):
    history_lbfgs.append(f(x_lbfgs).item())
    lbfgs.step(closure)


# Plotting
plt.semilogy(history_gd, label='GD')
plt.semilogy(history_lbfgs, label='L-BFGS')
plt.legend()
plt.show()

Correction: seems like calculating gradients for checking the second Wolfe-condition is necessary. I changed the gist. Thank you for pointing this out!

Thanks for sharing this. The docs are a bit unclear but I think if we are doing full dataset training (as opposed to minibatching) we can just set max_iter sufficiently high and then the for loop isn't needed (i.e., just call lbfgs.step once).

Thanks for this gist

Thanks for sharing this. The docs are a bit unclear but I think if we are doing full dataset training (as opposed to minibatching) we can just set max_iter sufficiently high and then the for loop isn't needed (i.e., just call lbfgs.step once).

I think so, but maybe lbfgs.step might need to be called twice to work. I'm using it on full dataset training, and running it just once doesn't seem to update anything. There's a couple if state['n_iter'] == 1: sections in lbfgs.step that initialize variables that are later computed when n_iter is greater than one.