使用 tensorflow-directml-plugin 为 TensorFlow 2 启用 GPU 加速

重要

此项目现已停用,不再积极开发。

此版本为学生、初学者和专业人士提供了一种使用 TensorFlow 2 的 DirectML 插件在其现有 DirectX 12 启用的硬件上运行机器学习(ML)培训的方法。

注释

可以使用 Python x86-64 3.10 进行安装 tensorflow-directml-plugin 。 但版本 3.11 及更高版本不支持 tensorflow-directml-plugin

了解如何将设备配置为使用 tensorflow-directml-pluginGPU 运行和训练模型。

步骤 1:最低(和最大)系统要求

在安装 TensorFlow-DirectML-Plugin 之前,请确保 Windows 或 WSL 版本支持 TensorFlow-DirectML-Plugin。

Windows 本机

  • Windows 10 版本 1709、64 位(内部版本 16299 或更高版本)或 Windows 11 版本 21H2、64 位(内部版本 22000 或更高版本)
  • Python x86-64 3.7、3.8、3.9 或 3.10。 版本 3.10 也是支持 的最大 版本。
  • 以下受支持的 GPU 之一:
    • AMD Radeon R5/R7/R9 2xx 系列或更高版本
    • Intel HD Graphics 5xx 或更高版本
    • NVIDIA GeForce GTX 9xx 系列 GPU 或更高版本

适用于 Linux 的 Windows 子系统

安装最新的 GPU 驱动程序

确保为硬件安装了最新的 GPU 驱动程序。 在“设置”应用的“Windows 更新”部分选择“检查更新”。 如果需要,请使用上述链接从硬件供应商处获取安装。

步骤 2:配置 Windows 环境

Windows 本机

本机 Windows 上的 TensorFlow-DirectML-Plugin 包从 Windows 10 版本 1709(内部版本 16299 或更高版本)开始工作。 可以查看内部版本号,方法是通过 Run 命令 (Windows logo key + R) 运行 winver

适用于 Linux 的 Windows 子系统

安装上述驱动程序后,请确保 启用 WSL安装基于 glibc 的分发版 (例如 Ubuntu 或 Debian)。 对于我们的测试,我们使用了 Ubuntu。 通过在“设置”应用的“Windows 更新”部分中选择“检查更新”,确保你拥有最新的内核。

注释

请确保启用更新 Windows 时接收其他 Microsoft 产品的更新选项。 可以在“设置”应用的 Windows 更新部分的“高级”选项中找到它。

对于这些功能,需要 5.10.43.3 或更高版本的内核版本。 可以在 PowerShell 中运行以下命令来检查版本号。

wsl cat /proc/version

步骤 3:设置环境

建议在 Windows 中设置虚拟 Python 环境。 有许多工具可用于设置虚拟 Python 环境,对于这些说明,我们将使用 Anaconda 的 Miniconda。 此设置的其余部分假定你使用 Miniconda 环境。 详细了解如何使用 Python 环境

在 Miniconda 中创建环境

在系统上下载并安装 Miniconda Windows 安装程序 。 Anaconda 网站上提供了关于设置的附加指南。 安装 Miniconda 后,使用名为 tfdml_plugin 的 Python 创建环境,并通过以下命令激活它。

conda create --name tfdml_plugin python=3.9 

conda activate tfdml_plugin 

注释

tensorflow 版本 >= 2.9,python 版本 >= 3.7 受支持。

步骤 4:安装基本 TensorFlow

下载基本 TensorFlow 包。 目前,directml 插件仅适用于 tensorflow–cpu==2.10,而不适用于 tensorflowtensorflow-gpu.

pip install tensorflow-cpu==2.10

步骤 5:安装 tensorflow-directml-plugin

安装此包会自动为现有脚本启用 DirectML 后端,而无需进行任何代码更改。

pip install tensorflow-directml-plugin

注释

如果训练脚本将设备字符串硬编码为非 "GPU",则可能会引发错误。

或者,可以从源生成包。 从源生成 tensorflow-directml-plugin 的说明

带有 DirectML 示例和反馈的 TensorFlow

查看 我们的示例 或使用你的现有模型脚本。 如果遇到问题,或者对 TensorFlow-DirectML-Plugin 包有反馈,请 与我们的团队联系