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

重要

此项目现已停用,目前尚未得到积极开发。

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

注意

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

了解如何使用 tensorflow-directml-plugin 配置设备以通过 GPU 运行和训练模型。

步骤 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 徽标键 + 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-plugin 仅适用于 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 的说明

TensorFlow with DirectML 示例和反馈

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