微软近日开源了 TensorFlow-DirectML 项目，它是 Tensorflow 的一个分支，利用 DirectML 在 Windows 10 和 WSL 上提供硬件加速。
DirectML 是微软计划为 WSL2 提供的两种 GPU 计算功能之一。它本身能够通过广泛支持的硬件和驱动程序为常见的机器学习任务提供 GPU 加速，包括来自 AMD、Intel、NVIDIA 和 Qualcomm 等供应商的所有支持 DirectX 12 的 GPU。
今年六月份，WSL2 已达到了对 GPU 的初始支持，其中就包括对 DirectML 的支持。当时，微软发布了带有 DirectML 后端的 TensorFlow 预览包，而这次开源的 TensorFlow-DirectML 是一个完整的 Tensorflow 分支。
![]()
微软高级程序经理 Clarke Rahrig 等人表示，TensorFlow-DirectML 扩展了 TensorFlow 的范围，DirectML 的加入使其超越了原有的 GPU 支持。
“机器学习工作负载通常涉及大量计算，尤其是在训练模型时。诸如 GPU 之类的专用硬件通常用于加速这些工作负载。TensorFlow 可以同时利用 CPU 和 GPU，但其 GPU 加速仅限于特定于供应商的平台，这些平台对 Windows 的支持及其用户的各种硬件范围都各不相同……而 DirectML 能够通过提供数学运算的高性能实现来扩展平台……这不仅扩展了 TensorFlow 在 Windows 上的 GPU 覆盖范围，而且还适用于 WSL，用户可以使用任何 DirectX 12 GPU 在 Windows 或 WSL 环境中运行或训练其 TensorFlow 模型。”
根据微软的介绍，DirectML 后端与 TensorFlow 集成，其新的内核集建立在 DirectML API 之上，而不是像现有的 CPU 和 GPU 内核一样基于 Eigen 源代码构建。二者的集成涉及到设备运行时（device runtime）的概念，设备运行时负责管理设备内存、在主机之间复制 tensors，记录 GPU 命令以及安排和同步在主机/CPU 和设备时间轴上发生的工作。下图显示了该设备运行时的一些关键组件，以及它与 DirectX 平台的接口方式：
![]()
- DmlDevice：实现 TensorFlow 的“设备”类，并最终管理所有与设备相关的功能
- DmlKernelWrapper：实现 TensorFlow 的“OpKernel”接口，该接口允许特定于设备的操作员实现
- DmlKernel：通过调用 DirectML 提供 TF 运算符的具体实现
- DmlKernelManager：缓存 DmlKernel 实例以避免在可能的情况下重新编译 DirectML 运算符
- DmlAllocator：管理支持 TensorFlow tensors 的 GPU 缓冲区
- DmlExecutionContext：调度在 GPU 上的工作，例如执行运算符或复制内存
TensorFlow-DirectML 在本机 Win32 和 WSL 上均可使用。更多详情可查阅微软的介绍文章。
