This document is relevant for: Inf2, Trn1, Trn1n

Select a Different Framework or Platform for Setup

TensorFlow Neuron (“tensorflow-neuronx”) Setup on Ubuntu 20 with DLAMI TensorFlow

Get Started with Latest Release of TensorFlow Neuron (tensorflow-neuronx)

This section provide links that will assist you to quickly start with a fresh installation of TensorFlow Setup Guide for Inf2 & Trn1 for both Inference and Training.

Launch the Instance

• Please follow the instructions at launch an Amazon EC2 Instance to launch an instance. When choosing the instance type at the EC2 console. please make sure to select the correct instance type.

• To get more information about instances sizes and pricing see: Trn1 web page, Inf2 web page

• Check for the latest version of the Deep Learning AMI Neuron TensorFlow 2.10 and copy the AMI name that starts with “Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 20.04) <latest_date>” from “AMI Name:” section

• Search for the copied AMI name in the AMI Search , you should see an exact matching AMI with the AMI name in Community AMIs. Select the AMI and use it to launch the instance.

• When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.

• After launching the instance, follow the instructions in Connect to your instance to connect to the instance

Update Neuron Drivers

# Update OS packages 
sudo apt-get update -y

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

# Install git 
sudo apt-get install git -y

# upgrade Neuron Driver
sudo apt-get upgrade aws-neuronx-dkms=2.* -y --allow-change-held-packages

# Update Neuron Runtime
sudo apt-get upgrade aws-neuronx-collectives=2.* -y --allow-change-held-packages
sudo apt-get upgrade aws-neuronx-runtime-lib=2.* -y --allow-change-held-packages

# Update Neuron Tools
sudo apt-get upgrade aws-neuronx-tools=2.* -y --allow-change-held-packages

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

Get Started With TensorFlow DLAMI

# Activate Python venv 
source /opt/aws_neuron_venv_tensorflow/bin/activate 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_tensorflow --display-name "Python (tensorflow-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Update Neuron Compiler and Framework
python -m pip install --upgrade neuronx-cc==2.* tensorflow-neuronx

Visit TensorFlow Neuron(tensorflow-neuronx) for Inference section

Visit TensorFlow Neuron section for more

Update to latest TensorFlow Neuron (tensorflow-neuronx)

If you already have a previous Neuron release installed, this section provide links that will assist you to update to latest Neuron release.

Tensorflow 2.10.1

Note

• Instructions in this page only apply to setting up Neuron components on Linux host running Ubuntu or Amazon Linux AMI.

• When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.

# Activate Python venv 
source /opt/aws_neuron_venv_tensorflow/bin/activate 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_tensorflow --display-name "Python (tensorflow-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Update Neuron Compiler and Framework
python -m pip install --upgrade neuronx-cc==2.* tensorflow-neuronx

Tensorflow 2.9.3

Note

• Instructions in this page only apply to setting up Neuron components on Linux host running Ubuntu or Amazon Linux AMI.

• When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.

# Activate Python venv 
source /opt/aws_neuron_venv_tensorflow/bin/activate 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_tensorflow --display-name "Python (tensorflow-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Update Neuron Compiler and Framework
python -m pip install --upgrade neuronx-cc==2.* tensorflow-neuronx==2.9.3.2.1.0

Tensorflow 2.8.4

Note

• Instructions in this page only apply to setting up Neuron components on Linux host running Ubuntu or Amazon Linux AMI.

• When launching a Trn1, please adjust your primary EBS volume size to a minimum of 512GB.

# Activate Python venv 
source /opt/aws_neuron_venv_tensorflow/bin/activate 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_tensorflow --display-name "Python (tensorflow-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Update Neuron Compiler and Framework
python -m pip install --upgrade neuronx-cc==2.* tensorflow-neuronx==2.8.4.2.1.0

Install Previous TensorFlow Neuron Releases for Ubuntu (tensorflow-neuronx)

This section will assist you to install previous Neuron releases.

Neuron 2.17.0

# Install Python venv 
sudo apt-get install -y python3.8-venv g++ 

# Create Python venv
python3.8 -m venv aws_neuron_venv_tensorflow 

# Activate Python venv 
source aws_neuron_venv_tensorflow/bin/activate 
python -m pip install -U pip 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_tensorflow --display-name "Python (tensorflow-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Install wget, awscli 
python -m pip install wget 
python -m pip install awscli 

# Install Neuron Compiler and Framework
python -m pip install neuronx-cc==2.12.68.0 tensorflow-neuronx==2.10.1.2.1.0

Neuron 2.16.0

# Install Python venv 
sudo apt-get install -y python3.8-venv g++ 

# Create Python venv
python3.8 -m venv aws_neuron_venv_tensorflow 

# Activate Python venv 
source aws_neuron_venv_tensorflow/bin/activate 
python -m pip install -U pip 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_tensorflow --display-name "Python (tensorflow-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Install wget, awscli 
python -m pip install wget 
python -m pip install awscli 

# Install Neuron Compiler and Framework
python -m pip install neuronx-cc==2.12.54.0 tensorflow-neuronx==2.10.1.2.1.0

Neuron 2.15.0

# Install Python venv 
sudo apt-get install -y python3.8-venv g++ 

# Create Python venv
python3.8 -m venv aws_neuron_venv_tensorflow 

# Activate Python venv 
source aws_neuron_venv_tensorflow/bin/activate 
python -m pip install -U pip 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_tensorflow --display-name "Python (tensorflow-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Install wget, awscli 
python -m pip install wget 
python -m pip install awscli 

# Install Neuron Compiler and Framework
python -m pip install neuronx-cc==2.11.0.34 tensorflow-neuronx==2.10.1.2.1.0

This document is relevant for: Inf2, Trn1, Trn1n