Last October we announced our generally consumable DirectML backend for TensorFlow 1.15. Since then, our team worked to bring machine learning training acceleration to TensorFlow 2. Today, we are happy to announce the release of our DirectML PluggableDevice preview package for TensorFlow on PyPI. Alongside our release, we are excited to announce that TensorFlow-DirectML-Plugin is open-sourced on GitHub.

What is TensorFlow-DirectML-Plugin

TensorFlow-DirectML-Plugin builds DirectML as a PluggableDevice backend to TensorFlow 2 for machine learning training on Windows and the Windows Subsystem for Linux. DirectML is an ML library that enables model acceleration across all DirectX 12 compatible GPUs.

Our pluggable device enables users of the latest version of TensorFlow to accelerate model training on a broad range of DX12-capable GPUs, including cards from AMD, Intel, and NVIDIA.

Using TensorFlow-DirectML-Plugin

Using TensorFlow-DirectML-Plugin for TensorFlow 2.9 is simple. Our pluggable device package is installable through PyPI without requiring any changes to your already-existing scripts.

The plugin works with TensorFlow core and easily integrates with versions 2.9 and newer of the tensorflow or tensorflow-cpu packages to seamlessly register your existing GPU.

Learn more about installing in our Docs

Try out the TensorFlow-DirectML-Plugin Today

We want to encourage all of you to pick up our TensorFlow-DirectML-Plugin and try it in your current workflow. If you prefer a tutorial, we have prepared samples for training SqueezeNet on GitHub.

This initial preview of our plugin package will support most basic machine learning models with increased model support and performance optimizations planned for subsequent releases.

Leave any questions, suggestions, or issues here on GitHub. Our team is constantly engaging with the community and would love to hear your input!