optiviz 1.0.4

Educational tool for visualising PyTorch-compatible ML optimisers on any differentiable 1D or 2D function.

[NEW] Walk along a loss landscape, watch from the sidelines, or ride along as your favourite optimisers battle it out for the best minimum with OptiViz 1.0, the latest update to OptiViz.

OptiViz

Optiviz is a python package that enables effortless visualisation of any PyTorch optimiser on any differentiable function in one or two variables. OptiViz might find educational use in an introductory nonlinear optimisation or deep learning class.

Usage

Installation

To install OptiViz, please use:

pip install optiviz

API

All functionality of OptiViz is exposed through the optiviz.optimise and optiviz.optimise_interactive functions.

import torch
from optiviz import optimise, optimise_interactive

Defining an objective

Before solving an optimisation problem you must define an objective function. OptiViz works with differentiable, real-valued objective functions in one or two variables.

f : \mathbb{R} \rightarrow \mathbb{R}

g : \mathbb{R}^2 \rightarrow \mathbb{R}

In code, you must define an objective function whose every input is a torch.Tensor of shape (1,)

def f(x: torch.Tensor) -> torch.Tensor:
    """
    Example of an objective function in one variable.
    """
    return x ** 2
def g(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Example of an objective function in two variables.
    """
    return x ** 2 + y ** 2 + x.sin() * y.sin()

optimise(): simple matplotlib visualisation

The optiviz.optimise function (please see docstring for complete usage) is used to visualise the optimisation sequence of a 1D or 2D objective function using a PyTorch optimiser.

arg_g_min = optimise(
        g, # objective function
        (12.5, 12.5), # initial values of the parameters being adjusted
        plot_centre=(0, 0), # plot options
        plot_boundary=25,
        iters=100,
        optimiser=torch.optim.Adam, # PyTorch-compatible optimiser
        lr=5e-1 # any keyword arguments for the optimiser
    )

optimise_interactive(): advanced interactive visualisation

The optiviz.optimise_interactive function (please see docstring for complete usage) provides interactive visualisation of 2D objective functions with multiple optimisers, navigation along the landscape, optimiser tracking, music and more in a web interface.

optimise_interactive(
        g, # objective function
        (12.5, 12.5), # initial values of the parameters being adjusted
        plot_centre=(0, 0), # plot options
        plot_boundary=25,
        iters=100,
        optimisers = [("Vanilla GD", lambda params: torch.optim.SGD(params, lr=0.1))], # PyTorch-compatible optimisers
        iter_delay: int = 100,
        gimbal_radius: float = 4.0, # advanced visualisation options
        gimbal_hover: float = 8.0,
    )

Example programmes

OptiViz 0.x/1.x:

import torch
from optiviz import optimise

f = lambda x,y: ×**2+y**2 # simple bowl

optimise(f, init_vector=(12.5, 12.5), optimiser=torch.optim.SGD, lr=le-2)

Optiviz 1.x:

import torch
from optiviz import optimise_interactive

def egg_carton(x, y): # non-convex function
    return 0.05*(x**2+y**2)+2.5*(torch.sin(0.5*x)**2+torch.sin(0.5*y)**2)

optimise_interactive(fn=egg_carton, optimisers=[("SGD with momentum", lambda params: torch.optim.SGD(params, lr=0.1, momentum=0.99)), ("Adam", lambda params: torch.optim.Adam(params, lr=0.5))])