TensorFlow and PyTorch on Graviton2 using Anaconda

Anaconda advertises 25 million users worldwide and offers an easy way to perform machine learning on a single machine.

Until recently, it has been difficult to install TensorFlow and PyTorch on Arm machines, including Graviton. There are numerous articles outlining difficult installation methods, most of which involve building source code. As the Arm software ecosystem continues to improve, more and more software just works on Arm. One of the latest improvements is Anaconda.

Today, let’s see how to get started with two popular machine learning frameworks, TensorFlow and PyTorch, on Graviton.

Setup a Graviton2 EC2 instance

First, create a new EC2 instance. The instance can be any of the instance types powered by Graviton processors including A1, T4g, M6g, C6g, or R6g. There are numerous AWS tutorials about how to create an AWS account and use the AWS Console to configure and launch a new EC2 instance. Make sure to substitute one of the above instance types in any tutorial you use.

The t4g.micro instances are free (up to 750 hours per month) until December 31, 2021. With 2 vCPUs and 1 GiB memory, the t4g.micro won't provide enough performance for machine learning projects, but does work to try out the steps in this article.

For this example, I created a t4g.medium running Ubuntu 20.04 with a public IP address and accessible from my desktop via ssh.

Find the public IP address of the EC2 instance in the AWS console and connect.

ssh -i mykey.pem ubuntu@<ec2-ip-address>

Once connected, update the software, install Anaconda, and source the bash environment.

The Anaconda install script will review the license agreement and ask to accept the terms. Then accept the default installation directory of $HOME/anconda3.

$ sudo apt update ; sudo apt upgrade -y 
$ wget https://repo.anaconda.com/archive/Anaconda3-2021.05-Linux-aarch64.sh
$ sh ./Anaconda3-2021.05-Linux-aarch64.sh
$ . ~/.bashrc

The default conda environment is named base and this will be shown in the shell prompt.

(base) ubuntu@ip-10-0-0-251:~$

Install TensorFlow

Create a new conda environment named tf, install TensorFlow, and activate the new environment.

$ conda create -n tf tensorflow
$ conda activate tf

The shell prompt will now show the tf environment.

(tf) ubuntu@ip-10-0-0-251:~$

Run a simple check to make sure TensorFlow is working.

$ python
Python 3.9.6 (default, Aug 19 2021, 13:29:06)
[GCC 10.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> print(tf.__version__)
2.5.0
>>> print(tf.reduce_sum(tf.random.normal([1000,1000])))
2021-08-27 17:21:47.972262: I tensorflow/core/common_runtime/process_util.cc:146] Creating new thread pool with default inter op setting: 2. Tune using inter_op_parallelism_threads for best performance.
tf.Tensor(-955.8242, shape=(), dtype=float32)
>>> exit()

Exit the tf environment and move on to install PyTorch.

$ conda deactivate

Install PyTorch

Create a new conda environment named torch, install PyTorch, and activate the new environment.

$ conda create -n torch pytorch
$ conda activate torch

The shell prompt will now show the tf environment.

(torch) ubuntu@ip-10-0-0-251:~$

Run a simple check to make sure TensorFlow is working.

$ python
Python 3.9.6 (default, Aug 19 2021, 13:29:06)
[GCC 10.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> x = torch.rand(5,3)
>>> print(x)
tensor([[1.1011e-01, 8.9852e-02, 2.9613e-01],
        [5.8129e-01, 3.6840e-01, 4.8840e-04],
        [2.8552e-01, 9.0869e-02, 3.5286e-01],
        [5.6122e-01, 9.2791e-01, 9.2109e-01],
        [7.4704e-01, 7.4170e-01, 8.9604e-01]])
>>> exit()

Try the MNIST example from PyTorch.

$ git clone https://github.com/pytorch/examples.git
$ pip install -r requirements.txt
$ cd examples
$ python main.py --epochs 2

There are many machine learning articles and examples using MNIST and now they can be done on Graviton2.

Wrap-up

I have spent many hours building machine learning frameworks from source for the Arm architecture. Although it’s challenging and satisfying to succeed, I can really appreciate how easy Anaconda makes getting started with TensorFlow and PyTorch on Graviton. Previously, the setup time was significant. With Anaconda, the setup time is short and there is more time available to create machine learning applications.

It's great to see an Anaconda installer just for Graviton2 on the bottom of the installation page.