Habilitación de PyTorch con DirectML en WSL

PyTorch con DirectML permite el entrenamiento y la inferencia en GPU compatibles con DirectX 12 en Subsistema de Windows para Linux (WSL). PyTorch con DirectML está en versión preliminar pública y funciona en WSL 2.

Comprobación de la versión de Windows

El torch-directml paquete de WSL 2 requiere Windows 11, compilación 22000 o posterior. Para comprobar la versión de Windows y el número de compilación, seleccione tecla del logotipo de Windows + R, escriba winver y seleccione Aceptar.

Instalación de WSL 2

Para instalar la distribución predeterminada de Linux con WSL 2, abra PowerShell o el Símbolo del sistema de Windows como administrador y ejecute:

wsl --install

Reinicie la máquina cuando se le solicite. Para obtener la selección de distribución y otras opciones de instalación, consulte Instalación de Linux en Windows con WSL.

Buscar actualizaciones de controladores de GPU

Instale el controlador de Windows más reciente disponible para la GPU a través de Windows Update o el sitio web del fabricante de hardware. El controlador Windows habilita la aceleración de GPU en WSL; no es necesario instalar un controlador de pantalla de Linux independiente.

Configuración de Python

Instale un entorno de Python en la distribución de WSL. Por ejemplo, ejecute los siguientes comandos para instalar Miniconda:

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

A continuación, cree y active un entorno denominado pytorch-directml:

conda create --name pytorch-directml python=3.10
conda activate pytorch-directml

Instalación de PyTorch con DirectML

Instala el paquete torch-directml:

pip install torch-directml

Comprobación de la instalación

Inicie Python y ejecute el código siguiente para agregar dos tensores en el dispositivo DirectML:

import torch
import torch_directml

dml = torch_directml.device()
tensor1 = torch.tensor([1]).to(dml)
tensor2 = torch.tensor([2]).to(dml)
dml_algebra = tensor1 + tensor2
print(dml_algebra.item())

Resultado esperado:

3

Ejemplos y comentarios

Consulte los ejemplos de PyTorch de DirectML para obtener ejemplos. Para notificar problemas de paquetes o solicitar características, use el seguimiento de problemas de DirectML.