PyTorch in Google Colab

Colab + PyTorch

UPDATE: Recently Google Colab comes preinstalled with PyTorch, which removes the need to install it manually, unless a specific version is required.

However, it is still useful to know how to install PyTorch in Google Colab environment. For example, PyTorch Hub functionality failed in the version installed in Google Colab and the temporary solution was to install the nighly build of PyTorch.

Based on the PyTorch website, we need a few parameters to have the correct version installation of PyTorch, which are CUDA version and Python version in our case.

How to check CUDA version in Colab?

!nvcc --version

NOTE: As of September 2019, the installed Cuda version on Google Colab is 10.

How to check Python version in Colab?

!python --version

Now we can easily install PyTorch using the parameter selector on PyTorch website:

Source: https://pytorch.org/

!pip3 install http://download.pytorch.org/whl/cu92/torch-0.4.1-cp36-cp36m-linux_x86_64.whl
!pip3 install torchvision

To confirm GPU is set up, you can run the following line after importing PyTorch:

import torch
torch.cuda.get_device_name(0)

My result in November 2018 was Tesla K80 GPU on Google Colab.

Source: https://www.nvidia.com/en-gb/data-center/tesla-k80/

Testing on September 2019, the GPU used in Google Colab is Tesla T4.

Source: https://www.nvidia.com/en-us/data-center/tesla-t4/

Tesla T4 specifications are as follows:

Without an actual announcement, the GPU on Google Colab is reported to be Tesla P100 for some people. The specs are:

Reference

Testing on October 2020, the GPU that I am getting is Tesla P100-PCIE-16GB.

Reference: NVIDIA

Reference: NVIDIA

Please take a moment to leave a comment or clap to let me know if this article was helpful.

What extra information would have been helpful? What else would you like to learn about?