【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 资源，从而加速模型训练。
如果你希望将这些方法应用到你的实际项目中，可以参考我的下一篇文章。如果你有任何问题或建议，欢迎在评论区留言，我会尽力帮助大家解决问题！