viz_torch_optim/main.ipynb

This project generates animations of pytorch optimizers solving toy problems.

Some nice animations were posted a few years ago by Alex Radford but didn't include Adam or landscapes with noise. Louis Tiao blogged about how to make the visulisations. The pytorch unit tests show how to run the optimizers on a test function. I pulled these together and shared the result at https://github.com/wassname/visualizing_torch_optim.

References:

• Alex Radford lost his original code but has something similar

• functions:

  • wikipedia test functions for opt
  • ssurjano a list of function and thier minima

• online symbolic algebra https://cocalc.com/settings?session=default

TODO:

• faster saving, to a format that can be converted well
• quality saving
• lr scheudle
• add noise
• fix float errors on complex functions

In [1]:

%pylab inline
from matplotlib import rcParams
rcParams['figure.figsize']=(10,10)

Populating the interactive namespace from numpy and matplotlib

In [2]:

import torch
from torch.autograd import Variable

In [ ]:

In [3]:

import matplotlib.pyplot as plt

from tqdm import tqdm_notebook as tqdm

from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import LogNorm
from matplotlib import animation
from IPython.display import HTML

from collections import defaultdict
from itertools import zip_longest
from functools import partial

In [4]:

import datetime
ts = datetime.datetime.utcnow().strftime('%Y%m%d_%H-%M-%S')

In [ ]:

In [5]:

from torch import optim
from torch.optim import lr_scheduler
torch.set_default_tensor_type('torch.DoubleTensor')
dtype=np.float64

Helpers to create optimization paths

In [ ]:

In [6]:

"""modified from https://github.com/pytorch/pytorch/blob/master/test/test_optim.py"""

def test_f(f, df, constructor, steps=150, x0=[-4,-1], solution=[-2,0], scheduler=None, exact=False):
    """
    params:
    scheduler: e.g. scheduler = torch.optim.CyclicLR(optimizer)
    """
    state = {}

    # start
    params = Variable(torch.Tensor(x0), requires_grad=True)
    optimizer = constructor([params])
    initial_lr = optimizer.param_groups[0]['lr']
    if scheduler:
        _scheduler = scheduler(optimizer, initial_lr)

    solution = torch.Tensor(solution)
    initial_dist = params.data.dist(solution)

    def eval():
        optimizer.zero_grad()
        loss = f(params)
        loss.backward()
        
        if exact:
            # loss.backward() will give **slightly** different
            # gradients, than drosenbtock, because of a different ordering
            # of floating point operations. In most cases it doesn't matter,
            # but some optimizers are so sensitive that they can temporarily
            # diverge up to 1e-4, just to converge again. This makes the
            # comparison more stable.
#             params.grad.data.copy_(df(params.data))
            pass
        return loss

    data=[]
    dist=[]
    lrs=[]
    for i in range(steps):
        
        loss = optimizer.step(eval)
        if scheduler:
            _scheduler.batch_step()
        
        # record
        dist.append(loss.squeeze().data.numpy()[0]) # loss
        data.append(params.data.numpy().copy())
        lrs.append(optimizer.param_groups[0]['lr'])
    return np.array(data), np.array(dist), lrs

In [7]:

# from https://github.com/thomasjpfan/pytorch/blob/401ec389db2c9d2978917a6e4d1101b20340d7e7/torch/optim/lr_scheduler.py
# not merged into pytorch yet https://github.com/pytorch/pytorch/pull/2016
from torch.optim import Optimizer
class CyclicLR(object):
    """Sets the learning rate of each parameter group according to
    cyclical learning rate policy (CLR). The policy cycles the learning
    rate between two boundaries with a constant frequency, as detailed in
    the paper `Cyclical Learning Rates for Training Neural Networks`_.
    The distance between the two boundaries can be scaled on a per-iteration
    or per-cycle basis.
    Cyclical learning rate policy changes the learning rate after every batch.
    `batch_step` should be called after a batch has been used for training.
    To resume training, save `last_batch_iteration` and use it to instantiate `CycleLR`.
    This class has three built-in policies, as put forth in the paper:
    "triangular":
        A basic triangular cycle w/ no amplitude scaling.
    "triangular2":
        A basic triangular cycle that scales initial amplitude by half each cycle.
    "exp_range":
        A cycle that scales initial amplitude by gamma**(cycle iterations) at each
        cycle iteration.
    This implementation was adapted from the github repo: `bckenstler/CLR`_
    Args:
        optimizer (Optimizer): Wrapped optimizer.
        base_lr (float or list): Initial learning rate which is the
            lower boundary in the cycle for eachparam groups.
            Default: 0.001
        max_lr (float or list): Upper boundaries in the cycle for
            each parameter group. Functionally,
            it defines the cycle amplitude (max_lr - base_lr).
            The lr at any cycle is the sum of base_lr
            and some scaling of the amplitude; therefore
            max_lr may not actually be reached depending on
            scaling function. Default: 0.006
        step_size (int): Number of training iterations per
            half cycle. Authors suggest setting step_size
            2-8 x training iterations in epoch. Default: 2000
        mode (str): One of {triangular, triangular2, exp_range}.
            Values correspond to policies detailed above.
            If scale_fn is not None, this argument is ignored.
            Default: 'triangular'
        gamma (float): Constant in 'exp_range' scaling function:
            gamma**(cycle iterations)
            Default: 1.0
        scale_fn (function): Custom scaling policy defined by a single
            argument lambda function, where
            0 <= scale_fn(x) <= 1 for all x >= 0.
            mode paramater is ignored
            Default: None
        scale_mode (str): {'cycle', 'iterations'}.
            Defines whether scale_fn is evaluated on
            cycle number or cycle iterations (training
            iterations since start of cycle).
            Default: 'cycle'
        last_batch_iteration (int): The index of the last batch. Default: -1
    Example:
        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
        >>> scheduler = torch.optim.CyclicLR(optimizer)
        >>> data_loader = torch.utils.data.DataLoader(...)
        >>> for epoch in range(10):
        >>>     for batch in data_loader:
        >>>         scheduler.batch_step()
        >>>         train_batch(...)
    .. _Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186
    .. _bckenstler/CLR: https://github.com/bckenstler/CLR
    """

    def __init__(self, optimizer, base_lr=1e-3, max_lr=6e-3,
                 step_size=2000, mode='triangular', gamma=1.,
                 scale_fn=None, scale_mode='cycle', last_batch_iteration=-1):

        if not isinstance(optimizer, Optimizer):
            raise TypeError('{} is not an Optimizer'.format(
                type(optimizer).__name__))
        self.optimizer = optimizer

        if isinstance(base_lr, list) or isinstance(base_lr, tuple):
            if len(base_lr) != len(optimizer.param_groups):
                raise ValueError("expected {} base_lr, got {}".format(
                    len(optimizer.param_groups), len(base_lr)))
            self.base_lrs = list(base_lr)
        else:
            self.base_lrs = [base_lr] * len(optimizer.param_groups)

        if isinstance(max_lr, list) or isinstance(max_lr, tuple):
            if len(max_lr) != len(optimizer.param_groups):
                raise ValueError("expected {} max_lr, got {}".format(
                    len(optimizer.param_groups), len(max_lr)))
            self.max_lrs = list(max_lr)
        else:
            self.max_lrs = [max_lr] * len(optimizer.param_groups)

        self.step_size = step_size

        if mode not in ['triangular', 'triangular2', 'exp_range'] \
                and scale_fn is None:
            raise ValueError('mode is invalid and scale_fn is None')

        self.mode = mode
        self.gamma = gamma

        if scale_fn is None:
            if self.mode == 'triangular':
                self.scale_fn = self._triangular_scale_fn
                self.scale_mode = 'cycle'
            elif self.mode == 'triangular2':
                self.scale_fn = self._triangular2_scale_fn
                self.scale_mode = 'cycle'
            elif self.mode == 'exp_range':
                self.scale_fn = self._exp_range_scale_fn
                self.scale_mode = 'iterations'
        else:
            self.scale_fn = scale_fn
            self.scale_mode = scale_mode

        self.batch_step(last_batch_iteration + 1)
        self.last_batch_iteration = last_batch_iteration

    def batch_step(self, batch_iteration=None):
        if batch_iteration is None:
            batch_iteration = self.last_batch_iteration + 1
        self.last_batch_iteration = batch_iteration
        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
            param_group['lr'] = lr

    def _triangular_scale_fn(self, x):
        return 1.

    def _triangular2_scale_fn(self, x):
        return 1 / (2. ** (x - 1))

    def _exp_range_scale_fn(self, x):
        return self.gamma**(x)

    def get_lr(self):
        step_size = float(self.step_size)
        cycle = np.floor(1 + self.last_batch_iteration / (2 * step_size))
        x = np.abs(self.last_batch_iteration / step_size - 2 * cycle + 1)

        lrs = []
        param_lrs = zip(self.optimizer.param_groups, self.base_lrs, self.max_lrs)
        for param_group, base_lr, max_lr in param_lrs:
            base_height = (max_lr - base_lr) * np.maximum(0, (1 - x))
            if self.scale_mode == 'cycle':
                lr = base_lr + base_height * self.scale_fn(cycle)
            else:
                lr = base_lr + base_height * self.scale_fn(self.last_batch_iteration)
            lrs.append(lr)
        return lrs

Functions

Tip when making new problems:

• get them from here

• and calculate the derivative using sagemath. Go to https://cocalc.com/ and execute code like this:

x = var('x')
y = var('y')
r=(1.5 - x + x*y)**2 + (2.25 - x + x*y**2)**2 + (2.625 - x + x*y**3)**2
print 'dx =',simplify(derivative(r,x))
print 'dy =',simplify(derivative(r,y))
# output:
# dx = 2*(x*y^3 - x + 2.625)*(y^3 - 1) + 2*(x*y^2 - x + 2.25)*(y^2 - 1) + 2*(x*y - x + 1.5)*(y - 1)
# dy = 6*(x*y^3 - x + 2.625)*x*y^2 + 4*(x*y^2 - x + 2.25)*x*y + 2*(x*y - x + 1.5)*x

In [8]:

azim=-95

In [9]:

def to_tensor(x):
    # TODO: I'm sure there's a proper way to do this
    if isinstance(x, np.ndarray):
        return Variable(torch.Tensor(x.astype(dtype)))
    if isinstance(x, list):
        return Variable(torch.Tensor(x))
    elif isinstance(x, (float, int, numpy.generic)):
        return Variable(torch.Tensor([float(x)]))
    else:
        return x
    
# def from_tensor(x):
#     x = getattr(x,'data',x)
#     if hasattr(x, 'numpy'):
#         x = x.numpy()
#     return x
# to_tensor(1.0)
# to_tensor(1)
# to_tensor(np.array([1.0]))
# to_tensor([1.0])
# to_tensor(torch.rand((4,4)))
# to_tensor(Variable(torch.rand((4,4))))

In [10]:

class Problem(object):
    def __init__(self, f, df, minima, x0, bounds=[[-5,5],[-5,5]], lr=1e-3, steps=3000, noise=dict(m=0,c=0)):
        """
        Problem setup
        
        Params:
        - f: function [x1,x2] => z
        - df: derivative function ([x1,x2]=>[dx1,dx2])
        - minima: where the function has a minima
        - self: bounds
        
        - x0: suggested start
        - lr: suggested learning rate
        - steps: suggested steps
        """
        def f_noise(t):
            """Add some noise"""
            t = to_tensor(t)
            c = Variable(torch.rand(t[0].size()) * noise['c'])
            m = (1 + Variable(torch.rand(t[0].size()) * noise['m']))
            z = f(t)
            return  m * z + c 
        self.f = f_noise
        self._f = f
        self.df = df
        self.x0 = x0
        self.bounds = bounds
        self.minima = minima
        self.lr = lr
        self.steps = steps
        
        self.xmin = bounds[0][0]
        self.xmax = bounds[0][1]
        self.ymin = bounds[1][0]
        self.ymax = bounds[1][1]
        
    

In [11]:

"""A valley"""
def madsen(tensor):
    """Madsen function (1981)."""
    x1, x2 = tensor
    x1 = to_tensor(x1)
    x2 = to_tensor(x2)

    r = x1**2 + x2**2 + x1 * x2 +\
            torch.sin(x1) +\
            torch.cos(x2) #+ np.abs(noise(x1,x2))
    return r
        
        
def dmadsen(tensor):
    x1, x2 = tensor
    x1 = to_tensor(x1)
    x2 = to_tensor(x2)
    dx1=2. * x1 + x2 + np.cos(x1)
    dx2 = 2. * x2 + x1 -np.sin(x2)
    return  torch.stack([dx1, dx2],1)[0]
 

madsen_problem= Problem(
    f=madsen,
    df=dmadsen,
    minima=np.array([-0.39999999999999591, 0.20000000000000462]),
    x0=[-3,-4],
    steps=3000,
    lr=1e-4
)
problem=madsen_problem

In [12]:

def schaffern4(tensor):
    """https://www.sfu.ca/~ssurjano/schaffer4.html"""
    x1, x2 = tensor
    x1 = to_tensor(x1)
    x2 = to_tensor(x2)

    r = 0.5 + (cos(sin(abs(x1**2-x2**2))) - 0.5)/(1+0.001*(x1**2+x2**2))**2
    return r*100


def dschaffern4(tensor):
    x1, x2 = tensor
    x1 = to_tensor(x1)
    x2 = to_tensor(x2)
    dx1 = -2*(x1**2 - x2**2)*x1*cos(abs(-x1**2 + x2**2))*sin(sin(abs(-x1**2 + x2**2)))/((0.001*x1**2 + 0.001*x2**2 + 1.0)**2*abs(-x1**2 + x2**2)) - 0.004*x1*(cos(sin(abs(-x1**2 + x2**2))) - 0.5)/(0.001*x1**2 + 0.001*x2**2 + 1.0)**3
    dx2 = 2*(x1**2 - x2**2)*x2*cos(abs(-x1**2 + x2**2))*sin(sin(abs(-x1**2 + x2**2)))/((0.001*x1**2 + 0.001*x2**2 + 1.0)**2*abs(-x1**2 + x2**2)) - 0.004*x2*(cos(sin(abs(-x1**2 + x2**2))) - 0.5)/(0.001*x1**2 + 0.001*x2**2 + 1.0)**3
    return torch.stack([dx1, dx2], 1)[0]


schaffern4_problem = Problem(
    f=schaffern4,
    df=dschaffern4,
#     minima=np.array([0, 1.25313]),
    minima=np.array([100,100]),
    bounds=[[-100, 100], [-100, 100]],
    x0=[10, 5],
    steps=3000,
    lr=1e-1
)
problem = schaffern4_problem

In [13]:

"""Banana shaped"""
def rosenbrock(tensor):
    x, y = tensor
    x = to_tensor(x)
    y = to_tensor(y)
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2


def drosenbrock(tensor):
    x, y = tensor
    x = to_tensor(x)
    y = to_tensor(y)
    dx = 400 * (x**2 -y) + 2 * x - 2 
    dy = -200 * x**2 + 200 * y
    return torch.stack([dx,dy],1)[0]


rosenbrock_problem= Problem(
    f=rosenbrock,
    df=drosenbrock,
    minima=np.array([1,1]),
    x0=[-3,-4],
    steps=3000,
    lr=1e-4
)
problem=rosenbrock_problem

In [14]:

from torch import sin, cos, sqrt, abs, log
def eggholder(tensor):
    """https://www.sfu.ca/~ssurjano/camel6.html"""
    x1, x2 = tensor
    x1 = to_tensor(x1)
    x2 = to_tensor(x2)

    r = - (x2+47)*sin(sqrt(abs(x2+x1/2+47))) - x1*sin(sqrt(abs(x1-(x2+47))))
    return r


def deggholder(tensor):
    x1, x2 = tensor
    x1 = to_tensor(x1)
    x2 = to_tensor(x2)
    dx1 = -1/8*(x1 + 2*x2 + 94)*(x2 + 47)*cos(sqrt(abs(1/2*x1 + x2 + 47)))/abs(1/2*x1 + x2 + 47)**(3/2) - 1/2*(x1 - x2 - 47)*x1*cos(sqrt(abs(-x1 + x2 + 47)))/abs(-x1 + x2 + 47)**(3/2) - sin(sqrt(abs(-x1 + x2 + 47)))
    dx2 = -1/4*(x1 + 2*x2 + 94)*(x2 + 47)*cos(sqrt(abs(1/2*x1 + x2 + 47)))/abs(1/2*x1 + x2 + 47)**(3/2) + 1/2*(x1 - x2 - 47)*x1*cos(sqrt(abs(-x1 + x2 + 47)))/abs(-x1 + x2 + 47)**(3/2) - sin(sqrt(abs(1/2*x1 + x2 + 47))) 
    return torch.stack([dx1, dx2], 1)[0]


eggholder_problem = Problem(
    f=eggholder,
    df=deggholder,
    minima=np.array([-512, -404.2319]),
    bounds=[[-512, 512], [-512, 512]],
    x0=[100,300],
    steps=3000,
    lr=1e-4
)
problem = eggholder_problem

In [15]:

def six_humped_camel_back(tensor):
    """https://www.sfu.ca/~ssurjano/camel6.html"""
    x1, x2 = tensor
    x1 = to_tensor(x1)
    x2 = to_tensor(x2)

    r = (4 - 2.1 * x1**2 + x1**4 / 3) * x1**2 + \
        x1 * x2 +\
        (-4 + 4 * x2**2) * x2**2
    return r


def dsix_humped_camel_back(tensor):
    x1, x2 = tensor
    x1 = to_tensor(x1)
    x2 = to_tensor(x2)
    dx1 = -0.333333333333333 * (-4.00000000000000 * x1**3 + 12.6000000000000 * x1) * x1**2 - 0.666666666666667 * \
        (-1.00000000000000 * x1**4 + 6.30000000000000 *
         x1**2 - 12.0000000000000) * x1 + x2
    dx2 = 8 * x2**3 + 8 * (x2**2 - 1) * x2 + x1
    return torch.stack([dx1, dx2], 1)[0]


camel6_problem = Problem(
    f=six_humped_camel_back,
    df=dsix_humped_camel_back,
    minima=np.array([0.0898,-0.7126]),
    bounds=[[-3, 3], [-2, 2]],
    x0=[1.3, -1.5],
    steps=3000,
    lr=1e-3
)
problem = camel6_problem

In [85]:

"""Valley"""


def beales(tensor):
    """Beales function, like a valley"""
    x, y = tensor
    x = to_tensor(x)
    y = to_tensor(y)
    # + noise(x,y)
    return (1.5 - x + x * y)**2 + (2.25 - x + x * y**2)**2 + (2.625 - x + x * y**3)**2


def dbeales(tensor):
    x, y = tensor
    x = to_tensor(x)
    y = to_tensor(y)
    dx = 2 * (x * y**3 - x + 2.625) * (y**3 - 1) + 2 * (x * y**2 -
                                                        x + 2.25) * (y**2 - 1) + 2 * (x * y - x + 1.5) * (y - 1)
    dy = 6 * (x * y**3 - x + 2.625) * x * y**2 + 4 * \
        (x * y**2 - x + 2.25) * x * y + 2 * (x * y - x + 1.5) * x
    return torch.stack([dx, dy], 1)[0]

beales_problem= Problem(
    f=beales,
    df=dbeales,
    minima=np.array([3., 0.5]),
    bounds=[[-4.5,4.5],[-4.5,4.5]],
    x0=[2,1.7],
    steps=3400,
    lr=3e-3,
    noise=dict(m=0.13, c=7),
)
problem=beales_problem

azim=-95

In [86]:

beales_problem.f(beales_problem.x0)

Out [86]:

Variable containing:
 162.6408
[torch.DoubleTensor of size 1]

In [87]:

# from torch import sin, cos, sqrt, abs, log, sinh, cosh
# from numpy import pi

# def levi13(tensor):
#     """https://www.sfu.ca/~ssurjano/levy13.html ."""
#     x1, x2 = tensor
#     x1 = to_tensor(x1)
#     x2 = to_tensor(x2)
#     # + noise(x,y)
#     return sin(3*pi*x1)**2 + (x1-1)**2 * (1+sin(3*pi*x2)**2) + (x2-1)**2*(1+sinh(2*pi*x2)**2)


# def dlevi13(tensor):
#     x1, x2 = tensor
#     x1 = to_tensor(x1)
#     x2 = to_tensor(x2)
#     pi = to_tensor(pi)
#     dx1 = 6*pi*cos(3*pi*x1)*sin(3*pi*x1) + 2*(sin(3*pi*x2)**2 + 1)*(x1 - 1)
#     dx2 = 6*pi*(x1 - 1)**2*cos(3*pi*x2)*sin(3*pi*x2) + 4*pi*(x2 - 1)**2*cosh(2*pi*x2)*sinh(2*pi*x2) + 2*(sinh(2*pi*x2)**2 + 1)*(x2 - 1)
#     return torch.stack([dx, dy], 1)[0]

# levi13_problem= Problem(
#     f=levi13,
#     df=dlevi13,
#     minima=np.array([0,0]),
#     bounds=[[-10,10],[-10,10]],
#     x0=[1.4,1.7],
#     steps=6000,
#     lr=1e-3
# )
# problem=levi13_problem

# azim=-95

In [ ]:

Grid

In [88]:

# define boundaries
# xmin, xmax, xstep = -5, 5, .05
# ymin, ymax, ystep = -5, 5, .05
# x0 = np.array([3., 4.])



xmin = problem.xmin
xmax = problem.xmax
ymin = problem.ymin
ymax = problem.ymax
ystep = xstep= (xmax-xmin)/200.0
zeps = 1.1e-0 # we don't want the minima to be actual zero or we wont get any lines shown on a log scale
z_min = problem.f(problem.minima).data.numpy()

# and x, y, z
x, y = np.meshgrid(np.arange(xmin, xmax + xstep, xstep), np.arange(ymin, ymax + ystep, ystep))
z = problem.f([x, y]).data.numpy() 

if z.min()<z_min[0]:
    print('WARN: your minima is not the true minima', z_min[0], z.min())
    z_min[0]=z.min()
    
z += -z_min[0] + zeps  # we shift everything up so the min is 1, so we can show on log scale
logzmax=np.log(z.max()-z.min()+zeps)
z.min(), z.max(), z_min[0]

Out [88]:

WARN: your minima is not the true minima 3.58007957647 0.0870856759143

(1.1000000000000001, 194897.71189181451, 0.087085675914340951)

In [ ]:

In [89]:

# reshape some vars
minima_ = problem.minima.reshape(-1, 1)

_x0 = np.array([problem.x0]).T

In [90]:

# assert z.min()>=z_min[0], 'your minina is wrong'
# assert (problem.df(problem.minima).data.numpy()<=1e-3).all(), 'gradient should be close to 0 at minima'

In [ ]:

In [91]:

ax = plt.gca()
cm=ax.contour(x, y, z, levels=np.logspace(0, logzmax//2, 55), norm=LogNorm(), cmap=plt.cm.jet, alpha=0.15)
plt.colorbar(cm)
ax.plot(*minima_, 'r*', markersize=10)
ax.plot(*problem.x0, 'r+', markersize=10)
plt.title('debug: grid')

plt.show()

Run

In [112]:

# you could use a lr scheduler too
scheduler = lambda optimizer, initial_lr: CyclicLR(optimizer, base_lr=1e-5, max_lr=initial_lr, step_size=100, mode='exp_range', gamma=0.9983)
scheduler = None

In [113]:

import torch.optim as optim
lr = problem.lr
constructors = dict(
    # Good high dimensional optimizers sometimes do poorly in low D spaces, so we will lower the LR on simple optimisers
    # need smaller lr's sometimes
    SGD= lambda params: optim.SGD(params, lr=lr/80),
    momentum = lambda params: optim.SGD(params, lr=lr/80, momentum=0.9, nesterov=False, dampening=0),
    momentum_dampen = lambda params: optim.SGD(params, lr=lr/80, momentum=0.9, nesterov=False, dampening=0.3),
    nesterov = lambda params: optim.SGD(params, lr=lr/80, momentum=0.9, nesterov=True, dampening=0),
    nesterov_decay = lambda params: optim.SGD(params, lr=lr/80, momentum=0.9, nesterov=True, weight_decay=1e-4, dampening=0),
    
    # need larger lr's sometimes
    Adadelta = lambda params: optim.Adadelta(params),
    Adagrad = lambda params: optim.Adagrad(params, lr=lr*20),
    
    # 
    Adamax = lambda params: optim.Adamax(params, lr=lr*20),
    RMSprop = lambda params: optim.RMSprop(params, lr=lr*10),
    Adam = lambda params: optim.Adam(params, lr=lr*10),
#     Adam_decay = lambda params:  optim.Adam(params, lr=lr*10, weight_decay=1e-9),
    
    # need to read about these, might not be comparable
#     ASGD = lambda params: optim.ASGD(params, lr=lr),
#     Rprop = lambda params: optim.Rprop(params, lr=lr),
#     LBFGS = lambda params: optim.LBFGS(params),
)

In [114]:

results = {}
distance = {}

for name, constructor in tqdm(constructors.items()):
    data, dist, lrs = test_f(problem.f, problem.df, constructor, x0=problem.x0, steps=problem.steps, scheduler=scheduler)
    
    results[name] = data
    distance[name] = dist
len(results)

# calc paths and elevation
# methods = constructors.keys()
# paths = np.array([path.T for path in results.values()]) # should be (2,N) each
# zpaths = np.array([[problem.f(torch.Tensor(p)).data.numpy()-z_min[0] + zeps for p in path.T] for path in paths])
# paths.shape, zpaths.shape

# calc paths and elevation
methods = constructors.keys()
paths = np.array([path.T for path in results.values()]) # should be (2,N) each
zpaths = np.array([distance[name] - z_min[0] + zeps for name in methods])
paths.shape, zpaths.shape

# DEBUG: check z's
for i, name in enumerate(results):
    zmax = zpaths[i][np.isfinite(zpaths[i])].max()
    print(name, zmax, '\t', np.isfinite(zmax).all(), '\t', zmax.max()>z[:,0].max())
    
# clip zpaths
zmax = z.max()
zpaths[np.isfinite(zpaths)==False]=zmax
zpaths = np.clip(zpaths, 0, zmax)
zpaths.shape

Out [114]:

[Data output - unsupported data type map[string]interface {} for mime type application/vnd.jupyter.widget-view+json]

SGD 159.745175336 	 True 	 False
momentum 179.057224365 	 True 	 False
momentum_dampen 176.793649996 	 True 	 False
nesterov 179.372084071 	 True 	 False
nesterov_decay 168.574150199 	 True 	 False
Adadelta 167.046837184 	 True 	 False
Adagrad 178.961374825 	 True 	 False
Adamax 165.557915528 	 True 	 False
RMSprop 171.633120474 	 True 	 False
Adam 162.48834522 	 True 	 False

(10, 3400)

In [ ]:

# # maybe I want to print full names?
# full_names = {}
# for name, constructor in constructors.items():
#     x0 = [1,2]
#     params = Variable(torch.Tensor(x0), requires_grad=True)
#     o = constructor([params])
#     lr = o.param_groups[0]['lr']
#     full_names[name]=f'{name} lr={lr:1.1e}'
#     print(full_names[name])

Preview plots

In [ ]:

In [115]:

# Preview plots
# static preview 2d to let you debug your steps and learning rate

# loss
for i, name in enumerate(results):
    plt.plot(np.abs(distance[name]), label=name)
plt.legend()
plt.title('loss (mae)')
plt.ylim(0,zpaths[:,0].mean()*1.3)
plt.show()


# Position
plt.figure(figsize=(12,12))
ax = plt.gca()
for name in results:
    plt.scatter(*results[name].T, label=name, s=1)
plt.legend()
plt.xlim(xmin,xmax)
plt.ylim(ymin,ymax)

cm=ax.contour(x, y, z, levels=np.logspace(0, logzmax//2, 35), norm=LogNorm(), cmap=plt.cm.jet, alpha=0.15)
# plt.colorbar(cm)
ax.plot(*minima_, 'r*', markersize=10)
ax.plot(*problem.x0, 'r+', markersize=10)
plt.title('debug: paths')

plt.show()

# lr
plt.plot(lrs)
# plt.yscale('log')
plt.title('learning rate')
plt.show()

In [ ]:

In [ ]:

In [ ]:

In [ ]:

In [116]:

# static preview 3d
fig = plt.figure(figsize=(8, 5))
ax = plt.axes(projection='3d', elev=50, azim=azim)

ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1, edgecolor='none', alpha=.25, cmap=plt.cm.jet)
ax.plot(*minima_, problem.f(minima_).data.numpy(), 'r*', markersize=10)
ax.plot(*_x0, problem.f(_x0).data.numpy(), 'r+', markersize=10)

ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$z$')

ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))

# anim = TrajectoryAnimation3D(*paths, zpaths=zpaths, labels=methods, ax=ax)
# quick plot to let you debug your steps and learning rate
ax = plt.gca()
for i, name in enumerate(results):
    ax.scatter3D(*results[name].T, zpaths[i], label=name, s=1)
plt.legend()
plt.xlim(xmin,xmax)
plt.ylim(ymin,ymax)

ax.legend(loc='upper right')

Out [116]:

<matplotlib.legend.Legend at 0x7ff0f53fde48>

Animation

In [117]:

from matplotlib import rcParams
rcParams['figure.figsize']=(10,10)
rcParams['figure.dpi']=100
rcParams['animation.writer']='ffmpeg' # faster than fmmpeg
rcParams['savefig.dpi']=180
rcParams['animation.codec']='h264'

# rcParams['savefig.bbox'] = 'tight'
seconds = 20
fps = 30 # too low and you will miss the fast moving ones
rcParams['animation.bitrate']=-1 #fps*1000//25 # ~1mb/s
cuttoff=problem.steps//1 # if we want to crop the data to X steps, judging by the loss plot

decimation = int(np.round(cuttoff/(seconds*fps))) or 1 # don't need to plot every step
# assert cuttoff>seconds*fps
decimation, cuttoff,seconds*fps, problem.steps

Out [117]:

(6, 3400, 600, 3400)

In [118]:

if scheduler:
    params = Variable(torch.Tensor([0,1]), requires_grad=True)
    optimizer = optim.Adam([params])
    scheduler_name = type(scheduler(optimizer, 0.1)).__name__
else:
    scheduler_name='None'

title=f'function: f{problem._f.__name__}'
if scheduler:
    title += f', scheduler: {scheduler_name}'
title

Out [118]:

'function: fbeales'

In [119]:

save_file = 'videos/{name:}_{scheduler}_{ts:}'.format(name=problem._f.__name__, ts=ts, scheduler=scheduler_name)
save_file

Out [119]:

'videos/beales_None_20171124_04-24-27'

In [120]:

# from http://louistiao.me/notes/visualizing-and-animating-optimization-algorithms-with-matplotlib/
class TrajectoryAnimation(animation.FuncAnimation):
    
    def __init__(self, *paths, labels=[], fig=None, ax=None, frames=None, 
                 interval=60, repeat_delay=5, blit=True, **kwargs):

        if fig is None:
            if ax is None:
                fig, ax = plt.subplots()
            else:
                fig = ax.get_figure()
        else:
            if ax is None:
                ax = fig.gca()

        self.fig = fig
        self.ax = ax
        
        self.paths = paths

        if frames is None:
            frames = max(path.shape[1] for path in paths)
  
        self.lines = [ax.plot([], [], label=label, lw=2)[0] 
                      for _, label in zip_longest(paths, labels)]
        self.points = [ax.plot([], [], 'o', color=line.get_color())[0] 
                       for line in self.lines]

        super(TrajectoryAnimation, self).__init__(fig, self.animate, init_func=self.init_anim,
                                                  frames=frames, interval=interval, blit=blit,
                                                  repeat_delay=repeat_delay, **kwargs)

    def init_anim(self):
        for line, point in zip(self.lines, self.points):
            line.set_data([], [])
            point.set_data([], [])
        return self.lines + self.points

    def animate(self, i):
        for line, point, path in zip(self.lines, self.points, self.paths):
            line.set_data(*path[::,:i])
            point.set_data(*path[::,i-1:i])
        return self.lines + self.points

In [ ]:

In [121]:

class TrajectoryAnimation3D(animation.FuncAnimation):
    
    def __init__(self, *paths, zpaths, labels=[], fig=None, ax=None, frames=None, 
                 interval=60, repeat_delay=5, blit=True, **kwargs):

        if fig is None:
            if ax is None:
                fig, ax = plt.subplots()
            else:
                fig = ax.get_figure()
        else:
            if ax is None:
                ax = fig.gca()

        self.fig = fig
        self.ax = ax
        
        self.paths = paths
        self.zpaths = zpaths
        
        if frames is None:
            frames = max(path.shape[1] for path in paths)
  
        self.lines = [ax.plot([], [], [], label=label, lw=2)[0] 
                      for _, label in zip_longest(paths, labels)]

        super(TrajectoryAnimation3D, self).__init__(fig, self.animate, init_func=self.init_anim,
                                                  frames=frames, interval=interval, blit=blit,
                                                  repeat_delay=repeat_delay, **kwargs)

    def init_anim(self):
        for line in self.lines:
            line.set_data([], [])
            line.set_3d_properties([])
        return self.lines

    def animate(self, i):
        for line, path, zpath in zip(self.lines, self.paths, self.zpaths):
            line.set_data(*path[::,:i])
            line.set_3d_properties(zpath[:i])
        return self.lines

In [ ]:

In [ ]:

In [ ]:

2d anim

hint on saving https://eli.thegreenplace.net/2016/drawing-animated-gifs-with-matplotlib/

In [122]:

title

Out [122]:

'function: fbeales'

In [123]:

fig, ax = plt.subplots(figsize=(12, 8))
# fig.set_tight_layout(True)
ax.contour(x, y, z, levels=np.logspace(0, logzmax//2, 35), norm=LogNorm(), cmap=plt.cm.jet, alpha=0.5)
ax.plot(*minima_, 'r*', markersize=10)
ax.plot(*problem.x0, 'r+', markersize=10)

ax.set_title('{} (github.com/wassname/viz_torch_optim)'.format(title))
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')

ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))

anim = TrajectoryAnimation(*paths[:,:,:cuttoff:decimation], labels=methods, ax=ax, interval=1000//fps)

ax.legend(loc='upper left')

Out [123]:

<matplotlib.legend.Legend at 0x7ff0f024d160>

In [124]:

# Query the figure's on-screen size and DPI. Note that when saving the figure to
# a file, we need to provide a DPI for that separately.
print('fig size: {0} DPI, size in inches {1}'.format(
    fig.get_dpi(), fig.get_size_inches()))

fig size: 100.0 DPI, size in inches [ 12.   8.]

In [125]:

import matplotlib.animation
import logging
_log = logging.getLogger(matplotlib.animation.__file__)
_log.setLevel(logging.DEBUG)

Saving, fps o 30, dpi 180 are ok, mp4 is faster than other formats

ffmpeg does poort quality gifs, use writer="imagemagick", for videos mp4 is quick, webm is nice too

In [126]:

anim.save(save_file+'_2d'+'.mp4', fps=fps, writer="ffmpeg", codec='h264')

In [127]:

# anim.save(save_file+'2d'+'.gif', fps=fps, codec='gif', writer="imagemagick") # slow, makes a large gif file

3d anim

In [128]:

fig = plt.figure(figsize=(12, 12) )
fig.set_tight_layout(True)
ax = plt.axes(projection='3d', elev=50, azim=azim)

ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1, edgecolor='none', alpha=.5, cmap=plt.cm.jet)
ax.plot(*minima_, problem.f(minima_).data.numpy(), 'r*', markersize=10)
ax.plot(*_x0, problem.f(_x0).data.numpy(), 'r+', markersize=10)

ax.set_title(title)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$z$')

ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))

anim = TrajectoryAnimation3D(*paths[:,:,:cuttoff:decimation], zpaths=zpaths[:,:][:,:cuttoff:decimation], labels=methods, ax=ax, interval=1000//fps)

ax.legend(loc='upper right')

Out [128]:

<matplotlib.legend.Legend at 0x7ff0cfd69400>

/home/isisilon/.pyenv/versions/3.6.0/envs/jupyter3/lib/python3.6/site-packages/matplotlib/figure.py:1743: UserWarning: This figure includes Axes that are not compatible with tight_layout, so its results might be incorrect.
  warnings.warn("This figure includes Axes that are not "

In [ ]:

In [129]:

# anim.save(save_file+'3d'+'.gif', fps=fps, codec='gif', writer="imagemagick")
anim.save(save_file+'_3d'+'.mp4', fps=fps, writer="ffmpeg", codec='h264')

loss anim

same as 2d but instead of x,y it's steps,np.abs(z-zmin)

In [ ]:

In [130]:

index = np.arange(1,problem.steps+1,1)[None,:]
index= np.array([index]*zpaths.shape[0])
loss_paths = np.concatenate([index, zpaths[:,None,:,0]], 1)
loss_paths.shape

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-130-2881e2b450d4> in <module>()
      1 index = np.arange(1,problem.steps+1,1)[None,:]
      2 index= np.array([index]*zpaths.shape[0])
----> 3 loss_paths = np.concatenate([index, zpaths[:,None,:,0]], 1)
      4 loss_paths.shape

IndexError: too many indices for array

In [ ]:

fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlabel('$step$')
ax.set_ylabel('$loss$')
ax.set_title(title)

ax.set_xlim((0, loss_paths[:,0,:].max()))
ax.set_ylim((0, loss_paths[:,1,:].max()))

anim = TrajectoryAnimation(*loss_paths[:,:,:cuttoff:decimation], labels=methods, ax=ax)

ax.legend(loc='upper right')

In [ ]:

# anim.save(save_file+'loss'+'.gif', fps=fps, codec='gif', writer="imagemagick")
anim.save(save_file+'_loss'+'.mp4', fps=fps, writer="ffmpeg", codec='h264')

In [ ]:

In [ ]:

In [ ]: