如何在 Ubuntu 和 Linux 中安装 TensorFlow-GPU

简介

TensorFlow-GPU 是针对 NVIDIA GPU 进行优化的 TensorFlow 版本，可以加速深度学习模型的训练。本文将介绍如何在 Ubuntu 和 Linux 中安装 TensorFlow-GPU。

步骤

1. 配置 NVIDIA CUDA 和 cuDNN

首先，您需要安装 NVIDIA CUDA 和 cuDNN。在 Ubuntu 18.04 中，您可以使用以下命令安装：

$ sudo apt-get install nvidia-cuda-toolkit

然后，您需要从 NVIDIA 网站下载 cuDNN 并安装。以下是安装 cuDNN 的步骤：

1. 访问 https://developer.nvidia.com/rdp/cudnn-download 页面
2. 选择您的下载选项，然后下载对应的 cuDNN 压缩包
3. 解压压缩包，例如：

$ tar -xzvf cudnn-11.1-linux-x64-v8.0.4.30.tgz

4. 复制 cuDNN 库到 CUDA 目录下：

$ sudo cp -P cuda/lib64/* /usr/local/cuda/lib64/
$ sudo cp  cuda/include/* /usr/local/cuda/include/

2. 创建 Python 虚拟环境

借助 Python 虚拟环境，可以避免不同 Python 应用程序之间的依赖冲突。您可以使用以下命令创建 Python 虚拟环境：

$ sudo apt-get install python3-venv
$ python3 -m venv myenv
$ source myenv/bin/activate

3. 安装 TensorFlow-GPU

接下来，您可以使用 pip 命令安装 TensorFlow-GPU：

$ pip install tensorflow-gpu

如果需要特定版本的 TensorFlow-GPU，您可以使用以下命令：

$ pip install tensorflow-gpu==2.5.0

4. 测试安装是否成功

您可以使用以下 Python 代码测试 TensorFlow-GPU 是否成功安装：

import tensorflow as tf

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

# Create a Tensor.
tensor = tf.constant('Hello, TensorFlow!')

# Print the Tensor.
print(tensor)

结论

安装 TensorFlow-GPU 并不难，只需要遵循以上四个步骤即可。当然，在安装 TensorFlow-GPU 之前，您需要确保已经正确地配置了 NVIDIA CUDA 和 cuDNN。如果无法正常安装，请查看相关错误日志并尝试解决。