This document is relevant for: Inf1, Inf2, Trn1, Trn1n

TensorFlow Neuron#

TensorFlow Neuron unlocks high-performance and cost-effective deep learning acceleration on AWS Trainium-based and Inferentia-based Amazon EC2 instances.

TensorFlow Neuron enables native TensorFlow models to be accelerated on Neuron devices, so you can use your existing framework application and get started easily with minimal code changes.

Tensorflow Neuron Setup
Install Drivers and Tools
# Configure Linux for Neuron repository updates
sudo tee /etc/yum.repos.d/neuron.repo > /dev/null <<EOF
[neuron]
name=Neuron YUM Repository
baseurl=https://yum.repos.neuron.amazonaws.com
enabled=1
metadata_expire=0
EOF
sudo rpm --import https://yum.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB

# Update OS packages 
sudo yum update -y

# Install OS headers 
sudo yum install kernel-devel-$(uname -r) kernel-headers-$(uname -r) -y

# Install git 
sudo yum install git -y

# install Neuron Driver
sudo yum install aws-neuronx-dkms-2.* -y

# Install Neuron Tools 
sudo yum install aws-neuronx-tools-2.* -y

# Add PATH
export PATH=/opt/aws/neuron/bin:$PATH

# Configure Linux for Neuron repository updates
. /etc/os-release
sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main
EOF
wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -

# Update OS packages 
sudo apt-get update -y

# Install OS headers 
sudo apt-get install linux-headers-$(uname -r) -y

# Install git 
sudo apt-get install git -y

# install Neuron Driver
sudo apt-get install aws-neuronx-dkms=2.* -y

# Install Neuron Tools 
sudo apt-get install aws-neuronx-tools=2.* -y

# Add PATH
export PATH=/opt/aws/neuron/bin:$PATH
Tutorials (tensorflow-neuronx)

Note

To use Jupyter Notebook see:

API Reference Guide (tensorflow-neuronx)
Misc (tensorflow-neuronx)
Tutorials (tensorflow-neuron)

  • Tensorflow 1.x - Running TensorFlow BERT-Large with AWS Neuron [html]

  • Tensorflow 2.x - HuggingFace Pipelines distilBERT with Tensorflow2 Neuron [html] [notebook]

  • Tensorflow 1.x & 2.x - Using NEURON_RT_VISIBLE_CORES with TensorFlow Serving [html]

Note

To use Jupyter Notebook see:

Additional Examples (tensorflow-neuron)
API Reference Guide (tensorflow-neuron)
Misc (tensorflow-neuron)

Note

TensorFlow Neuron support for Training is coming soon.

This document is relevant for: Inf1, Inf2, Trn1, Trn1n