Lutz RoederDec 27, 2016

TensorFlow on Azure

Training neural networks (deep learning) is very compute-intensive. Fast GPUs can make those sessions, which sometimes take hours, days or weeks go orders of magnitude faster. However, laptops usually don't come with the fastest GPUs and having to maintain a desktop machine only to occasionally run deep learning tasks is extra hassle.

Cloud providers now offer virtual machines (VMs) with GPUs which run in data centers and can be used by anybody on an hourly basis. Below is a quick tutorial that walks through setting up a VM in Microsoft Azure with the necessary drivers to train neural networks using TensorFlow.

First, if you haven't done so already, create an Azure account, install the Azure 2.0 command line interface (CLI)...

sudo pip install azure-cli

... and follow the login procedure:

az login

Azure manages resources (virtual machines, storage etc.) via resource groups. GPU virtual machine instances are currently available in the East US region. If you already have a group for that region feel free to use it, otherwise create a new resource group:

az group create --name tensorflow --location EastUS

We will connect to the machine via SSH and need to create a key pair:

ssh-keygen -f ~/.ssh/tensorflow_id_rsa -t rsa -b 2048 -C '' -N ''

Next, we create the actual virtual machine running Ubuntu 16.04. We choose the cheapest and least powerful GPU size (NC6) and downgrade from premium (SSD) to standard storage (HDD) as the former is not supported for NC instances yet.

az vm create --resource-group tensorflow --name tensorflow --image Canonical:UbuntuServer:16.04-LTS:latest --size Standard_NC6 --storage-type Standard_LRS --admin-username tensorflow --ssh-key-value ~/.ssh/

Once completed, the command will print the IP address for the newly created machine:

  "publicIpAddress": "",
  "resourceGroup": "tensorflow"

The VM is now running in a data center (and charging for cycles). The following commands can be used to deallocate and restart anytime:

az vm deallocate --resource-group tensorflow --name tensorflow
az vm start --resource-group tensorflow --name tensorflow

Connect to the machine via SSH using the IP address printed above (type 'yes', if asked to continue):

ssh tensorflow@ -i ~/.ssh/tensorflow_id_rsa

Next, download CUDA, make it known to apt-get and run install:

sudo dpkg -i cuda-repo-ubuntu1604-8-0-local_8.0.44-1_amd64-deb
sudo apt-get update
sudo apt-get install -y cuda
rm cuda-repo-ubuntu1604-8-0-local_8.0.44-1_amd64-deb

Now we can check the status of the GPU(s) by running nvidia-smi.

Next, download and install cuDNN...

sudo tar -xzf cudnn-8.0-linux-x64-v5.1.tgz -C /usr/local
rm cudnn-8.0-linux-x64-v5.1.tgz
sudo ldconfig

...and add the following exports to ~/.bashrc:

export CUDA_HOME=/usr/local/cuda-8.0
export PATH=/usr/local/cuda-8.0/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-8.0/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

The final step is to install Pip and the GPU version of TensorFlow:

sudo apt-get install -y python-pip python-dev
sudo pip install -y tensorflow-gpu

We can now start a Python console and create a TensorFlow session:

>>> import tensorflow as tf
>>> session = tf.Session()

If everything went well, it will recognize the Tesla K80 GPU:

I tensorflow/core/common_runtime/gpu/]
Found device 0 with properties: 
name: Tesla K80
major: 3 minor: 7 memoryClockRate (GHz) 0.8235
pciBusID b0b5:00:00.0
Total memory: 11.17GiB
Free memory: 11.11GiB

Remember to deallocate the VM when done to avoid using cycles:

az vm deallocate --resource-group tensorflow --name tensorflow

Once no longer needed, you can delete the virtual machine by running:

az vm delete --resource-group tensorflow --name tensorflow