PyTorch

PyTorch project is a Python package that provides GPU accelerated tensor computation and high level functionalities for building deep learning networks. For licensing details, see the PyTorch license doc on GitHub.

To monitor and debug your PyTorch models, consider using TensorBoard.

The following sections provide guidance on installing PyTorch on Databricks and give an example of running PyTorch programs.

Note

This is not a comprehensive guide to PyTorch. Refer to the PyTorch website.

Install PyTorch

Databricks Runtime for ML

Databricks Runtime for Machine Learning includes PyTorch so you can create the cluster and start using PyTorch. Here are the Pytorch versions included:

Databricks Runtime ML Version PyTorch Version
Databricks Runtime 7.3 ML 1.6.0
Databricks Runtime 7.2 ML 1.5.1
Databricks Runtime 7.1 ML 1.5.1
Databricks Runtime 7.0 ML 1.5.0
Databricks Runtime 6.6 ML 1.4.0
Databricks Runtime 6.5 ML 1.4.0
Databricks Runtime 6.4 ML 1.4.0
Databricks Runtime 6.3 ML (Unsupported) 1.3.1
Databricks Runtime 5.5 LTS ML 1.1.0

Databricks Runtime

We recommend using the PyTorch included on Databricks Runtime for Machine Learning. However, if you must use Databricks Runtime, PyTorch can be installed as a Databricks PyPI library. The following example shows how to install PyTorch 1.5.0:

  • On GPU clusters, install pytorch and torchvision by specifying the following:
    • torch==1.5.0
    • torchvision==0.6.0
  • On CPU clusters, install pytorch and torchvision by using the following wheel files:
    • https://download.pytorch.org/whl/cpu/torch-1.5.0%2Bcpu-cp37-cp37m-linux_x86_64.whl
    • https://download.pytorch.org/whl/cpu/torchvision-0.6.0%2Bcpu-cp37-cp37m-linux_x86_64.whl

Use PyTorch on a single node

To test and migrate single-machine PyTorch workflows, you can start with a driver-only cluster on Databricks by setting the number of workers to zero. Though Apache Spark is not functional under this setting, it is a cost-effective way to run single-machine PyTorch workflows.

PyTorch notebook

Open notebook in new tab