【ShuQiHere】小白也能懂的 TensorFlow 和 PyTorch GPU 配置教程

【ShuQiHere】

在深度学习中，GPU 的使用对于加速模型训练至关重要。然而，对于许多刚刚入门的小白来说，如何在 TensorFlow 和 PyTorch 中指定使用 GPU 进行训练可能会感到困惑。在本文中，我将详细介绍如何在这两个主流的深度学习框架中指定使用 GPU 进行训练，并确保每一个步骤都简单易懂，跟着我的步骤来，你也能轻松上手！

1. 安装所需库

首先，确保你已经安装了 TensorFlow 或 PyTorch。你可以使用以下命令进行安装：

1. 安装 TensorFlow

pip install tensorflow

# 注意：对于 Linux 用户，如果你使用pip install tensorflow安装，只会安装CPU 版本
# 这是因为在 Linux 中，TensorFlow 的 GPU 版本和 CPU 版本是分开的。
# 为了在 GPU 上运行 TensorFlow，请使用下面的命令
pip install tensorflow[and cuda]

2. 安装 PyTorch

访问 PyTorch 的官网 并根据你的系统和 CUDA 版本选择合适的安装命令，例如：

pip install torch torchvision torchaudio

确保安装的 TensorFlow 和 PyTorch 版本支持你的 GPU。

2. 检查 GPU 可用性 (For Linux Users)

在进行训练之前，必须确保系统中有可用的 GPU，并且 TensorFlow 或 PyTorch 可以正确识别它们。

1. 使用 nvidia-smi 命令检查 GPU 状态

在终端中输入以下命令：

nvidia-smi

你会看到类似下面的输出：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 535.183.01   Driver Version: 535.183.01   CUDA Version: 12.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce RTX 3060   Off  | 00000000:01:00.0 Off |                  N/A |
| N/A   50C    P3    N/A /  55W |    8MiB /  6144MiB |      0%      Default   |
+-------------------------------+----------------------+----------------------+

这个命令可以帮助你确认 GPU 是否被正确识别，同时显示 GPU 的使用情况。

如果你无法识别到GPU，可以参考我的文章，先在Linux环境下安装并且配置GPU、CUDA、cuDNN。

【ShuQiHere】在Ubuntu上安装NVIDIA驱动-CSDN博客【ShuQiHere】在Ubuntu上安装CUDA和cuDNN-CSDN博客

2. 在 TensorFlow 和 PyTorch 中检查 GPU 可用性

在 TensorFlow 中，你可以使用以下代码检查 GPU 的可用性：

import tensorflow as tf

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

在 PyTorch 中，你可以使用以下代码检查：

import torch

print("Is CUDA available: ", torch.cuda.is_available())

如果输出结果显示存在可用的 GPU，那么你可以放心继续后面的步骤了。

3. 在训练中指定 GPU

在 TensorFlow 和 PyTorch 中，你可以分别通过以下方式指定使用 GPU 进行训练。

3.1 在 TensorFlow 中指定 GPU

1. 自动检测和使用 GPU

TensorFlow 会自动检测并使用可用的 GPU。你可以通过以下代码验证：

import tensorflow as tf

# 创建一个简单的操作
a = tf.random.normal([1000, 1000])
b = tf.random.normal([1000, 1000])
c = tf.matmul(a, b)

# 打印操作所在的设备
print(f"Operation done on device: {c.device}")

2. 手动指定 GPU

有时你可能需要手动指定使用哪个 GPU。以下是一个示例：

import tensorflow as tf

# 指定使用第一个 GPU
with tf.device('/GPU:0'):
    model = tf.keras.models.Sequential()
    # 添加网络层，定义模型结构
print("Model is set to run on GPU: /GPU:0")

3.2 在 PyTorch 中指定 GPU

1. 自动检测和使用 GPU

PyTorch 也会自动检测可用的 GPU。你可以通过以下代码确认：

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

2. 手动指定 GPU

在 PyTorch 中，你可以通过 torch.device 手动指定使用哪个 GPU：

import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = YourModel().to(device)
print("Model moved to:", device)

记得将所有的张量和模型都移动到指定的设备上：

tensor = torch.randn(3, 3).to(device)
print("Tensor moved to:", device)

4. 常见问题与故障排除

1. GPU 未被识别的问题

有时Linux环境下，TensorFlow 或 PyTorch 中其中一个库可能无法识别你的 GPU。常见原因可能是：

• CUDA 和 cuDNN 版本不匹配

• 驱动程序未更新

• 环境变量配置错误

可以通过检查 CUDA、cuDNN 和驱动程序的版本是否匹配来解决。

2. 内存不足的错误

如果你在训练过程中遇到“内存不足”的错误，可以尝试以下方法：

• 降低批次大小（batch size）

• 使用混合精度训练（mixed precision training）

5. 总结

在这篇文章中，我们详细介绍了如何在 TensorFlow 和 PyTorch 中指定使用 GPU 进行训练。从安装必要的库，到检查 GPU 可用性，再到指定 GPU 进行训练，我们进行了详细的讲解。通过这些步骤，你可以确保在深度学习任务中充分利用 GPU 资源，从而加速模型训练。

如果你希望将这些方法应用到你的实际项目中，可以参考我的下一篇文章。如果你有任何问题或建议，欢迎在评论区留言，我会尽力帮助大家解决问题！