Visualization in tutorial not working in colab and remote setup

[SebastianRiechert]

I have tried running this tutorial both in google colab and in jupyterlab on a remote server setup.

https://ax.dev/tutorials/tune_cnn.html

Both times the render()-call returns an empty output.

This is the code, no changes from tutorial except installation of ax:

!pip install ax-platform
%%
import torch
import numpy as np

from ax.plot.contour import plot_contour
from ax.plot.trace import optimization_trace_single_method
from ax.service.managed_loop import optimize
from ax.utils.notebook.plotting import render, init_notebook_plotting
from ax.utils.tutorials.cnn_utils import load_mnist, train, evaluate, CNN

init_notebook_plotting()
%%
torch.manual_seed(12345)
dtype = torch.float
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
%%
BATCH_SIZE = 512
train_loader, valid_loader, test_loader = load_mnist(batch_size=BATCH_SIZE)
%%
def train_evaluate(parameterization):
    net = CNN()
    net = train(net=net, train_loader=train_loader, parameters=parameterization, dtype=dtype, device=device)
    return evaluate(
        net=net,
        data_loader=valid_loader,
        dtype=dtype,
        device=device,
    )
%%
best_parameters, values, experiment, model = optimize(
    parameters=[
        {"name": "lr", "type": "range", "bounds": [1e-6, 0.4], "log_scale": True},
        {"name": "momentum", "type": "range", "bounds": [0.0, 1.0]},
    ],
    evaluation_function=train_evaluate,
    objective_name='accuracy',
)
%%
render(plot_contour(model=model, param_x='lr', param_y='momentum', metric_name='accuracy'))

switching to classic jupyter notebook (in remote setup) as suggested in https://github.com/facebook/Ax/issues/94 doesn't change anything.

pip freezes for both environments:

colab_pip_freeze.txt

jupyterlab_pip_freeze.txt

[adamobeng]

Thanks for flagging this, @SebastianRiechert! I've asked @2timesjay, who's most familiar with this part of the code, to look into it.

[2timesjay]

Hi @SebastianRiechert , the discussion on #332 lead to the solution for the jupyterlab half of this issue. copying that here:

This is unfortunately a problem with all JupyterLab plotly plotting.

I was able to resolve it by following https://plotly.com/python/getting-started/#jupyterlab-support-python-35 and installing the jupyter labextension suggested.

Please let us know if this works for you.

Alternately, Jupyter Notebooks should not encounter this problem.

The colab issue is unexpected: https://stackoverflow.com/questions/47230817/plotly-notebook-mode-with-google-colaboratory describes the issue as resulting from using older plotly versions, but we're on plotly >4.4 and still running into it. While I investigate further, the workaround from the top answer there DOES work in colab.

[lena-kashtelyan]

Closing as duplicate, let's direct all future discussion of this issue to #332.

[2timesjay]

Still investigating colab half of this issue (realized this one was the right one to host discussion on).

[lena-kashtelyan]

Support for Colab is considered wishlist for now, porting this to https://github.com/facebook/Ax/issues/566.

[lena-kashtelyan]

Workaround for making Ax plots work in colab is a function like this (source: https://stackoverflow.com/questions/47230817/plotly-notebook-mode-with-google-colaboratory, this seems to apply to Plotly verstions 3.x overall):

def configure_plotly_browser_state():
  import IPython
  display(IPython.core.display.HTML('''
        <script src="/static/components/requirejs/require.js"></script>
        <script>
          requirejs.config({
            paths: {
              base: '/static/base',
              plotly: 'https://cdn.plot.ly/plotly-latest.min.js?noext',
            },
          });
        </script>
        '''))

So you would call that function before rendering a plot, e.g.:

configure_plotly_browser_state()
render(ax_client.get_contour_plot())

We'll incorporate this into a convenient utility soon as well (maybe just making it part of render), at which point we'll close the issue.

[saitcakmak]

Here's a simple workaround. Add

import plotly.io as pio
pio.renderers.default = "colab"

to the NB running on colab. The figure renders and is interactive :) (for smaller notebook file size at the expense of non-interactive plots, you can use "png" as the renderer)

Note: You also need to do !pip install kaleido.

[lena-kashtelyan]

This is now covered in documentation: https://ax.dev/tutorials/visualizations.html#Fix-for-plots-that-are-not-rendering and a warning will be raised in the code directing users to the documentation. @septfreur, @sgbaird, @SebastianRiechert FYI