How To Set Up Google Colab/Colaboratory For Building Pyro Models

Google’s Colab is basically Google Docs for interactive, runnable Jupyter Notebooks, and it is a fantastic resource for playing around with machine learning models. It comes with great sharing and collaboration features and even a free GPU runtime!

However, the backend it runs on only comes with a subset of the dependencies needed to run Pyro models. I’m going to quickly run through the steps to set up Pyro in a Colab notebook. If you already know how to use Colab and just want the setup code for Pyro, skip to Install and import dependencies.

Sign up for Colab

Anyone with a Google Drive account can sign up for Colab by heading to colab.research.google.com and following the listed instructions.

Create a new notebook

File -> New Notebook. Either Python 2 or 3 will work, but I would recommend Python 3 (that’s what the rest of this tutorial will be using)

Set up the GPU runtime

Runtime -> Change runtime type -> Hardware Accelerator -> GPU. Free GPU cycles! 😍

Install and import dependencies

In the first code cell of your notebook, install PyTorch and Pyro. Those exclamation points tell Jupyter/Colab that the pip3 commands are meant to be executed as shell commands and not Python code.

Test it out

Try running the basic Pyro model from the SVI Part I tutorial:

You should get some output about the fairness of a coin’s flips.

That’s that! Happy deep universal probabilistic programming!

The full code for a working Pyro model with Python 3 on Colab can be found here.