使用 WSL 中用于ML的 GPU 加速开始

机器学习 (ML) 正在成为许多开发工作流的关键部分。 无论你是数据科学家、ML工程师,还是使用 ML 适用于 Linux 的 Windows 子系统 (WSL) 开始学习之旅,都提供了一个出色的环境,可以运行最常见的常用 GPU 加速ML工具。

设置这些工具的方法有很多不同。 例如,WSL、TensorFlow-DirectMLPyTorch-DirectML中的 NVIDIA CUDA 都提供了不同的方式,可以将 GPU 用于 WSL ML。 若要详细了解选择一个与另一个的原因,请参阅 GPU 加速ML训练

本指南将介绍如何设置:

  • 如果有 NVIDIA 图形卡并运行示例ML框架容器,则为 NVIDIA CUDA
  • 在 AMD、Intel 或 NVIDIA 图形卡上TensorFlow-DirectML和PyTorch-DirectML

先决条件

使用 Docker 设置 NVIDIA CUDA

  1. 下载并安装 NVIDIA GPU 的最新驱动程序

  2. 运行以下命令,在 WSL 中安装 Docker Desktop 或直接安装 Docker 引擎

    curl https://get.docker.com | sh
    
  3. 如果直接安装 Docker 引擎,请按照以下步骤安装 NVIDIA 容器Toolkit

    运行以下命令,为 NVIDIA 容器Toolkit设置稳定存储库:

    distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
    
    curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
    
    curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
    

    通过运行以下命令安装 NVIDIA 运行时包和依赖项:

    sudo apt-get update
    
    sudo apt-get install -y nvidia-docker2
    
  4. 运行机器学习框架容器和示例。

    若要运行机器学习框架容器并开始使用此 NVIDIA NGC TensorFlow 容器使用 GPU,请输入以下命令:

    docker run --gpus all -it --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/tensorflow:20.03-tf2-py3
    

    TensorFlow with CUDA running inside a Docker container

    可以通过运行命令来运行内置在此容器中的预先训练的模型示例:

    cd nvidia-examples/cnn/
    
    python resnet.py --batch_size=64
    

    TensorFlow sample model training within Docker container

有关获取设置和使用 NVIDIA CUDA 的其他方法,请参阅 WSL 用户指南上的 NVIDIA CUDA

设置TensorFlow-DirectML或PyTorch-DirectML

  1. 从 GPU 供应商网站下载并安装最新的驱动程序: AMDIntelNVIDIA

  2. 设置 Python 环境。

    建议设置虚拟 Python 环境。 有许多工具可用于设置虚拟 Python 环境 - 对于这些说明,我们将使用 Anaconda 的 Miniconda

    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
    
    bash Miniconda3-latest-Linux-x86_64.sh
    
    conda create --name directml python=3.7 -y
    
    conda activate directml
    
  3. 安装由所选 DirectML 支持的机器学习框架。

    TensorFlow-DirectML:

    pip install tensorflow-directml
    

    PyTorch-DirectML:

    sudo apt install libblas3 libomp5 liblapack3
    
    pip install pytorch-directml
    
  4. TensorFlow-DirectMLPyTorch-DirectML 的交互式 Python 会话中运行快速添加示例,以确保一切正常工作。

如果有疑问或遇到问题,请访问 GitHub 上的 DirectML 存储库

其他资源