在 WSL 2 上通过 DirectML 启用 PyTorch

2025-04-30

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

检查你的 Windows 版本

适用于 Linux 的 Windows 子系统 (WSL) 2 中的 torch-directml 包从 Windows 11（内部版本 22000 或更高版本）开始应用。可以查看内部版本号，方法是通过 Run 命令 (Windows logo key + R) 运行 winver。

检查 GPU 驱动程序更新

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

设置 Torch-DirectML

安装 WSL 2

若要安装适用于 Linux 的 Windows 子系统（WSL） 2，请参阅安装 WSL 中的说明。

然后按照 README.md GitHub 存储库中的文件中的说明安装 WSL GUI 驱动程序。

设置 Python 环境

建议在 WSL 2 中设置虚拟 Python 环境。有许多工具可用于设置虚拟 Python 环境，本主题将使用 Anaconda 的 Miniconda。此设置的其余部分假定你使用 Miniconda 环境。

按照 Anaconda 站点上的 Linux 安装程序指南或在 WSL 2 中运行以下命令来安装 Miniconda。

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh 
bash Miniconda3-latest-Linux-x86_64.sh

安装 Miniconda 后，创建名为 pytdml 的 Python 环境，并通过以下命令激活它：

conda create --name pytdml -y
conda activate pytdml

安装 PyTorch 和 Torch-DirectML

注释

torch-directml 包最多支持 PyTorch 2.3.1

要完成安装，只需通过运行以下命令来安装 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

包含 DirectML 示例和反馈的 PyTorch

请查看我们的示例，了解将 PyTorch 与 DirectML 配合使用的更多用法。如果遇到问题，或者有有关 PyTorch 的 DirectML 包的反馈，请在此处与我们的团队联系。

通过