如何设置 Python 使用 GPU

如果你想加速你的 Python 程序，使用 GPU 可以是一个不错的选择。同时，一些深度学习框架，如 TensorFlow 和 PyTorch，也需要 GPU 支持来实现更好的训练和预测性能。接下来，我们将介绍如何设置 Python 使用 GPU，以下步骤主要基于 NVIDIA 的 GPU。

步骤一：安装 NVIDIA 显卡驱动

首先，需要安装合适版本的 NVIDIA 显卡驱动。你可以在 NVIDIA 官网 根据你的显卡型号以及操作系统下载对应驱动。安装时，请参考官方文档或网上教程，确保安装过程无误。

步骤二：安装 CUDA

CUDA（Compute Unified Device Architecture）是 NVIDIA 提供的并行计算平台和编程模型，是使用 GPU 进行科学计算和深度学习任务的必备组件。你可以在 NVIDIA 官网 中根据你的显卡和操作系统下载对应的 CUDA 版本。安装时，请参考官方文档或网上教程，确保安装过程无误。

# 安装过程参考
# 下载 CUDA 安装包
wget https://developer.download.nvidia.com/compute/cuda/11.2.2/local_installers/cuda_11.2.2_460.32.03_linux.run

# 安装 CUDA
sudo sh cuda_11.2.2_460.32.03_linux.run

步骤三：安装 cuDNN

cuDNN（CUDA Deep Neural Network library）是一个针对深度学习应用的 GPU 加速库，支持一些主流深度学习框架。你可以在 NVIDIA 官网 中下载对应的 cuDNN 版本。注意，下载时需要先注册一个账号。

下载完毕后，将压缩包解压到 CUDA 的安装目录下即可。

# 解压 cuDNN 压缩包
tar -xzvf cudnn-11.2-linux-x64-v8.1.0.77.tgz

# 复制文件到 CUDA 目录下
sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h
sudo chmod a+r /usr/local/cuda/lib64/libcudnn*

步骤四：安装 Python 相关的 GPU 库

上述三步完成后，还需要安装一些 Python GPU 库，如 TensorFlow-gpu、PyTorch 等。你可以使用 pip 命令来安装它们。注意，需要选择与 CUDA 版本和 cuDNN 版本兼容的版本，否则可能会出现错误。

# 安装 TensorFlow-gpu
pip install tensorflow-gpu==2.4.0

# 安装 PyTorch
pip install torch torchvision torchaudio

好了，以上就是设置 Python 使用 GPU 的全部步骤。接下来，你可以尝试使用 GPU 运行 Python 程序，并享受 GPU 加速带来的快感。

参考资料：

• Installing TensorFlow with GPU support
• Getting Started with PyTorch and Deep Learning