Microsoft AI Open-Sources ‘PyTorch-DirectML’: A Package To Train Machine Learning Models On GPUs


The Microsoft Windows AI team has announced the first preview of DirectML as a backend to PyTorch for training ML models. This release allows accelerated machine learning training for PyTorch on any DirectX12 GPU and WSL, unlocking new potential in computing with mixed reality.

Microsoft AI team has teamed up with the PyTorch framework to release a preview package that provides scoped support for CNNs (convolutional neural networks). In this new device named “DML,” Direct ML APIs and Tensor primitives are called through by introducing minimal overhead when calling into operators; they work in much like other existing backends.

The PyTorch-DirectML suite can use the GPU machine learning library DirectML as its back end, allowing models to be trained on DirectX12 GPUs and WSL (Windows Subsystem for Linux).

Microsoft collaborated with AMD, Intel, and NVIDIA to provide this hardware-accelerated training experience for PyTorch. The PyTorch-DirectML package is an easy installation that only requires changing one line of code in an existing script.