AWS Machine Learning Blog

Updated AWS Deep Learning AMIs: New Versions of TensorFlow, Apache MXNet, Keras, and PyTorch

We’re excited to update the AWS Deep Learning AMIs with significantly faster training on NVIDIA Tesla V100 “Volta” GPUs across many frameworks, including TensorFlow, PyTorch, Keras, and the latest Apache MXNet 1.0 release.

There are two main flavors of the AMIs available today. The Conda-based AWS Deep Learning AMI packages the latest point releases of popular frameworks (along with pre-configured CUDA, notebook and supporting libraries), inside Conda-based virtual environments. It’s recommended for all developers. The Deep Learning AMI with Source Code is the right choice if you want to make changes to the underlying frameworks, or are comfortable on the bleeding edge. We make custom, optimized builds available on this AMI. You’ll find a lot of advanced features and performance improvements here; test your code first before moving into production with these. 

Faster training with mixed precision for TensorFlow

The new AMIs with Source Code include a customized, updated build of TensorFlow which takes advantage of the mixed-precision training and inference on V100 GPUs available in P3 instances on EC2. The result is significantly faster training times. For example, training a ResNet-50 using this build performs 1.6x better compared to stock TensorFlow 1.4.

This build is available on the updated AMIs with Source Code with CUDA 9 for Amazon Linux and Ubuntu; it’s based on the master branch of TensorFlow (commit: d73e8b3, pulled on 5th December, 2017). If you would prefer to use TensorFlow 1.4, it’s available in a Conda virtual environment on the AWS Deep Learning AMI for Amazon Linux and Ubuntu. 

Performance improvements for Keras 2.0

Keras 2.0 is now available in a Conda virtual environment on the AMIs, running on a TensorFlow 1.4 backend. Additionally, you can also run your Keras code against the Volta-optimized TensorFlow build described above on the Source Code AMIs with CUDA 9.

Latest versions of PyTorch, with Volta support

The Conda and Source Code AMIs now include the latest version of PyTorch, v0.3.0, which includes support for CUDA 9, cuDNN 7 and associated performance improvements for training models on V100 “Volta” GPUs. There are also a number of new performance and usability improvements, including a model exporter for ONNX – a common model interchange format supported by MXNet, Caffe2, and Microsoft Cognitive Toolkit.

Latest version of Apache MXNet 1.0, with performance, usability, and interoperability

The AMIs also include the Apache MXNet 1.0 release. A new model serving capability for MXNet packages, runs, and serves deep learning models with just a few lines of code. It includes a new gradient compression capability, and a new model converter that converts neural network code written with the Caffe framework to MXNet code, making it easier for developers to take advantage of MXNet’s scalability and performance.

Deep learning and ONNX

The Open Neural Network Exchange (ONXX), is an initiative supported by Facebook, Microsoft and AWS, to improve portability of trained models between frameworks. Many of the frameworks included in the Deep Learning AMIs now include some level of ONNX support, which makes the AMIs a great place to test the inference performance of your models on multiple engines, all without leaving your instance. For example, you can train in PyTorch, and infer in Caffe2 – the Deep Learning AMIs make this transition quick and painless, since you have all the latest frameworks and tools to make and evaluate this transition at your fingertips. 

Getting started with the AWS Deep Learning AMI

It’s fast and simple to get started with the AWS Deep Learning AMIs. The latest releases of the AMIs are available in the AWS Marketplace to support the various needs of machine learning practitioners. Our AMI selection guide helps you pick the right AMI for your deep learning project and gets you going with just one click. We’ve also provided many quick tutorials and developer resources to help you accelerate model training.

Conda-based AMI:

Deep Learning AMI (Ubuntu)

Deep Learning AMI (Amazon Linux)

AMI with source code:

For P3 instances

Deep Learning AMI with Source Code (CUDA 9, Ubuntu)

Deep Learning AMI with Source Code (CUDA 9, Amazon Linux)

For P2 instances

Deep Learning AMI with Source Code (CUDA 8, Ubuntu)

Deep Learning AMI with Source (CUDA 8, Amazon Linux)

 

 

About the Author

Cynthya Peranandam is a Principal Marketing Manager for AWS artificial intelligence solutions, helping customers use deep learning to provide business value. In her spare time she likes to run and listen to music.