在 Windows 上通过 DirectML 启用 PyTorch

PyTorch with DirectML 为开发人员提供一种易于使用的方式,可在其 Windows 计算机上试用最新的和最出色的 AI 模型。 可以通过安装 torch-directml PyPi 包下载 PyTorch with DirectML。 设置之后,即可开始使用我们的示例,或者使用用于 VS Code 的 AI 工具包。

检查 Windows 版本

本机 Windows 上的 torch-directml 包从 Windows 10 版本 1709(内部版本 16299 或更高版本)开始运行。 可以查看内部版本号,方法是通过 Run 命令(Windows 徽标键 + R)运行 winver

检查 GPU 驱动程序更新

确保已安装最新的 GPU 驱动程序。 在“设置”应用的“Windows 更新”部分中,选择“检查是否有更新”

设置 Torch-DirectML

建议在 Windows 内设置虚拟 Python 环境。 可以使用许多工具来设置虚拟 Python 环境 - 对于这些说明,我们将使用 Anaconda 的 Miniconda。 此设置的其余部分假定你使用 Miniconda 环境。

设置 Python 环境

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

conda create --name pytdml -y
conda activate pytdml

安装 PyTorch 和 Torch-DirectML

注意

torch-directml 包支持的最高版本为 PyTorch 2.2

获取安装所需的所有操作是通过运行以下命令来安装最新版本的 torch-directml

pip install torch-directml

验证和设备创建

安装 torch-directml 包后,可以通过添加两个张量来验证其是否正确运行。 首先启动交互式 Python 会话,并使用以下行导入 Torch:

import torch
import torch_directml
dml = torch_directml.device()

torch-directml 的当前版本映射到“PrivateUse1”Torch 后端。 torch_directml.device() API 是一个方便的包装器,用于将张量发送到 DirectML 设备。

创建 DirectML 设备后,现在可以定义两个简单的张量:一个张量包含 1,另一个张量包含 2。 将张量放在“dml”设备上。

tensor1 = torch.tensor([1]).to(dml) # Note that dml is a variable, not a string!
tensor2 = torch.tensor([2]).to(dml)

将张量相加,然后打印结果。

dml_algebra = tensor1 + tensor2
dml_algebra.item()

应该显示输出数字 3,如下例所示。

>>> import torch
>>> tensor1 = torch.tensor([1]).to(dml)
>>> tensor2 = torch.tensor([2]).to(dml)
>>> dml_algebra = tensor1 + tensor2
>>> dml_algebra.item()
3

PyTorch with DirectML 示例和反馈

请查看我们的示例,以了解 PyTorch with DirectML 的更多用法。 如果遇到问题或有关于 PyTorch with DirectML 包的反馈,请与我们的团队联系