成功地使用本地的 NVIDIA GPU 运行 PyTorch 或 TensorFlow

我用Windows 11的电脑编程，主要是用Python。有时候，我发现我电脑上的Python代码无法使用我安装的NVIDIA GPU。这是因为为了让像PyTorch或TensorFlow这样的Python库能利用你的GPU，你需要在电脑上安装一套完整的兼容软件。

在这篇文章里，我将分享您在Windows 11 PC上需要安装哪些内容及其顺序，以让Python程序能利用您现有的图形处理单元（GPU）。

要在系统上使用CUDA，您需要安装以下软件：

• 支持CUDA的GPU
• 支持gcc编译器和工具链的受支持版本的Linux
• NVIDIA CUDA工具包（可在https://developer.nvidia.com/cuda-downloads下载）

以下支持的 Microsoft Windows 操作系统版本：

• Microsoft Windows 11 21H2,
• Microsoft Windows 11 22H2-SV2,
• Microsoft Windows 11 23H2,
• Microsoft Windows 10 22H2,
• Microsoft Windows Server 2022

1. 要在Windows上使用CUDA工具包，首先需要安装的是Visual Studio。 这是因为CUDA编译使用的编译器是与Visual Studio捆绑的Microsoft Visual C++。CUDA还用于调试和性能分析，促进开发。安装Visual Studio的详细说明。
2. 安装完Visual Studio后，接下来需要安装NVIDIA CUDA Toolkit。

CUDA® 是由 NVIDIA 发明的，并行计算平台及编程模型。它能大幅提升计算性能，利用图形处理单元（GPU）的计算能力。

如果你的电脑配备了NVIDIA GPU，很可能已经安装好了CUDA。如果没有，你可以在这里下载它。完整的安装指南在这里查看

当你安装CUDA工具包时，会在你的程序文件夹中出现一个名为“NVIDIA GPU计算工具包文件夹”的文件夹。

在这个文件夹内会有子文件夹，当你进入名为“CUDA”的子文件夹时，你会看到lib、bin、include等子文件夹。

3. 安装NVIDIA图形驱动程序。可以参考此链接https://docs.nvidia.com/deeplearning/cudnn/latest/installation/windows.html。

4. 安装完上述要求后，还需要安装cuDNN库。

NVIDIA的CUDA深度神经网络库（cuDNN）是一个GPU加速的库，用于深度神经网络，提供了高效实现，包括前向和后向卷积、注意力、矩阵乘法、池化和归一化等常用操作。

你可以在这里下载cuDNN。

要下载cuDNN

• 访问 NVIDIA cuDNN 页面并使用您的 NVIDIA 开发者账户登录。如果您还没有账户，您需要创建一个。
• 前往 cuDNN 下载页面并下载与您的 CUDA 版本相匹配的版本（对于 CUDA 11.5，您应该下载相应的 cuDNN 版本）。

• 解压下载的文件。您会发现 bin、include 和 lib 等文件夹。将这些文件夹中的内容复制到您的 CUDA 安装路径中的相应目录（例如：C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5）：

• 将 cudnn\bin 目录中的文件复制到 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5\bin

• 将 cudnn\include 目录中的文件复制到 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5\include

• 将 cudnn\lib\x64 目录中的文件复制到 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5\lib\x64

• 检查并验证安装：复制完文件后，目录中应该包含必需的 cuDNN 文件：

• cudnn.h 应在 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5\include 目录中

• cudnn.lib 应在 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5\lib\x64 目录中

• 您可以运行 TensorFlow 或 PyTorch 等支持 GPU 的深度学习程序进行验证。

• 完成上述步骤后，您的系统现在应该已经安装好了 CUDA 和 cuDNN，可以执行 GPU 加速的任务。

这里是安装cuDNN的详细步骤详细步骤。你可以在这里查看你的CUDA Toolkit版本支持的cuDNN版本：支持矩阵。

在安装cuDNN之前，您必须先从这里安装适合您系统的NVIDIA图形驱动程序此处。这里。**

安装了NVIDIA图形驱动后，您可以选择图形安装或Tarball安装，在Windows上安装cuDNN。

一旦安装 cuDNN 后，您会在“Program Files”目录下的 NVIDIA GPU Computing Toolkit 目录中的 bin、include 和 lib\x64 文件夹中找到这些文件。

如果你有 bin、include 和 lib\x64 文件夹但没有找到 cudnn.h 或 cudnn.lib 文件，这表明，你的系统上没有安装 cuDNN。CUDA 和 cuDNN 是分开安装的，所以安装了 CUDA 并不意味着 cuDNN 也一定安装了。

5. 现在你只需要安装合适的 PyTorch 或者 TensorFlow 版本，以便利用你的 CUDA GPU。

你可以通过这个网站帮助你找到合适的pip或conda命令来安装适合你的CUDA环境的PyTorch安装版本。

来源：https://pytorch.org/get-started/locally/

这里是如何为GPU安装TensorFlow的说明。

6. 确认您的Python库能识别并利用GPU。

用于检查你的 PyTorch 是否能识别 GPU 的 Python 代码

这里是我在本地GPU上运行的完整代码，使用基于Stable Diffusion的生成式AI模型（stable-diffusion-xl-base-1.0）根据指定的提示生成图像。该模型名为“stable-diffusion-xl-base-1.0”，它来自Stability.AI在Hugging Face上的模型页面：stable-diffusion-xl-base-1.0。

    import torch  
    from diffusers import DiffusionPipeline  

    from huggingface_hub import notebook_login  
    notebook_login()  

    # 加载预训练模型
    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")  
    pipe.to("cuda")  

    prompt = “一个宇航员骑着一匹绿色的马”  
    images = pipe(prompt=prompt).images[0]  

    from PIL import Image  
    from IPython.display import display  
    display(images)

上述代码的图像输出。