We are happy to announce monthly releases of TensorFlow and PyTorch container images for Arm servers, with the first release being 21.04 in April 2021. With these images, end users can now easily run machine learning inference workloads on Arm servers like AWS Graviton2. In this article, let us explore the details of the 21.04 release and the plans for the upcoming releases. For detailed instructions on how to use the docker images, please refer to the blog on AArch64 docker images for TensorFlow and PyTorch.
The first release 21.04 has TensorFlow and PyTorch images with two different backends – Default and oneDNN + Compute Library for Arm Architecture (ACL). The machine learning frameworks use different backends depending on what is available in the hardware. TensorFlow, by default, uses the Eigen GEBP kernel on CPU-only systems while PyTorch, by default, uses a combination of libraries like OpenBLAS, QNNPACK for CPU-only systems. To improve inference performance, Arm is working towards a oneDNN+Compute Library backend, that allows frameworks to benefit from optimized machine learning primitives available in Compute Library for Arm Architecture (ACL).
oneAPI Deep Neural Network Library (oneDNN) is an open-source cross-platform performance library of basic building blocks for deep learning applications. Formerly known as MKL-DNN, the library has been integrated with TensorFlow, PyTorch, and other frameworks. Compute Library (ACL) is an Arm optimized library with over 100 optimized machine learning functions including multiple convolution algorithms. Arm has integrated ACL with oneDNN so that frameworks can run faster on Arm cores.
The 21.04 release ships TensorFlow 2.3.0 version with following backend choices:
The default Eigen backend is well tested and ready for production use. The oneDNN+ACL backend provides higher performance as you can see in the below graph and is currently under development.
We plan to make the oneDNN+ACL the default option on CPU-only Arm servers. We would love to hear from you on your experience with the oneDNN+ACL backend and your feedback helps us in identifying areas where functionality and performance need to be improved.
In addition to Arm-provided docker images, Linaro hosts TensorFlow Python packages for AArch64 and we are working with the TensorFlow development team to provide official AArch64 TensorFlow packages.
The 21.04 release ships PyTorch 1.8.1 version with backend choices:
The default backend is well tested and ready for production use. The official PyTorch packages for AArch64 available at https://download.pytorch.org/whl/torch_stable.html use the default backend. We are in the process of adding oneDNN+ACL backend support to PyTorch and the current release is a preview of what will be available in the future PyTorch packages.
Going forward, We at Arm aim to provide a monthly update of TensorFlow and Pytorch containers. We update the framework versions to newer releases, as and when they become available. We also plan to provide the improvements in oneDNN+ACL in the docker releases ahead of them being available upstream. We look forward to hearing about your experience with the new docker releases.
Get started