DirectML acceleration for PyTorch is currently available for Public Preview. PyTorch with DirectML enables training and inference of complex machine learning models on a wide range of DirectX 12-compatible hardware.

DirectML is a high-performance, hardware-accelerated DirectX 12 library for machine learning. DirectML provides GPU acceleration for common machine learning tasks across a broad range of supported hardware and drivers, including all DirectX 12-capable GPUs from vendors such as AMD, Intel, NVIDIA, and Qualcomm.

More information about DirectML can be found in Introduction to DirectML.

PyTorch on DirectML is supported on both the latest versions of Windows 10 and the Windows Subsystem for Linux, and is available for download as a PyPI package. For more information about getting started, see GPU accelerated ML training (docs.microsoft.com)

For users of Pytorch-DirectML forked from Pytorch 1.13 or higher, the samples can be found below or in the 1.13 folder:

• attenion is all you need- the original transformer model
• yolov3- a real-time object detection model
• squeezenet - a small image classification model
• resnet50 - an image classification model
• maskrcnn - an object detection model

For users of Pytorch-DirectML forked from Pytorch 1.8, the samples can be found below or in the 1.8 folder:

• squeezenet - a small image classification model
• resnet50 - an image classification model
• maskrcnn - an object detection model

• PyTorch homepage
• torch-directml PyPI project