Skip to content

TensorFlow 安装

在这篇文章中

TensorFlow 是由 Google 开发的用于机器学习和深度神经网络训练的开源平台。借助 TensorFlow，您可以使用 Python 中用户友好的 API 来创建、训练和利用复杂的神经网络。

该平台既支持传统的机器学习算法（如线性回归和逻辑回归），也支持更高级的架构，包括卷积神经网络（Convolutional Neural Networks, CNN）和循环神经网络（Recurrent Neural Networks, RNN）。

Linux 上的 TensorFlow 安装

本指南适用于以下操作系统：Ubuntu 22.04，并针对以下 Python 版本进行了验证：Python 3.10。

备注

如果您计划使用 GPU 加速，请根据 此指南 安装 NVIDIA 驱动程序和 CUDA。

  1. 安装 Python：

    sudo apt install python3.10

    在 Ubuntu 22.04 中，此版本默认已安装，因此我们不建议安装更新的版本。

  2. 安装 TensorRT 的系统库：

    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get install tensorrt

  3. 创建 Python 虚拟环境：

    python3 -m venv venv

  4. 激活虚拟环境：

    source venv/bin/activate

    成功激活后，提示符中将包含括号内的虚拟环境名称：

    (venv) user@49069:~$

    备注

    您可以创建任意数量的虚拟环境并安装不同的库（包括同时安装，但有时这可能会导致冲突）。

  5. 安装 TensorRT 绑定：

    python3 -m pip install wheel
python3 -m pip install --pre --upgrade tensorrt

    为了验证安装，我们可以在控制台运行 Python 并执行以下脚本：

    import tensorrt
print(tensorrt.__version__)
assert tensorrt.Builder(tensorrt.Logger())

    安装成功后，输出将如下所示：

    (tensorflow) user@49069:~/gpu$ python
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorrt
>>> print(tensorrt.__version__)
10.0.0b6
>>> assert tensorrt.Builder(tensorrt.Logger())
>>> import tensorrt_lean as trt
>>> print(trt.__version__)
10.0.0b6
>>> assert trt.Runtime(trt.Logger())
>>> import tensorrt_dispatch as trt
>>> print(trt.__version__)
10.0.0b6
>>> assert trt.Runtime(trt.Logger())
>>>

  6. 安装 TensorRT 版本 8.6.1。

    wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.1/tars/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-12.0.tar.gz
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-12.0.tar.gz

  7. 安装带有 GPU 支持的 TensorFlow 库：

    pip install tensorflow[and-cuda]

  8. 退出虚拟环境：

    deactivate

  9. 创建运行 TensorFlow 的脚本：

    echo '#!/bin/bash' > tensorflow.sh
echo 'source venv/bin/activate' >> tensorflow.sh
echo 'export CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))' >> tensorflow.sh
echo 'export LD_LIBRARY_PATH=$CUDNN_PATH/lib:/usr/local/cuda/lib64' >> tensorflow.sh
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/TensorRT-8.6.1.6/lib' >> tensorflow.sh
chmod +x tensorflow.sh

运行 TensorFlow

要在主目录的根目录下使用指定的库变量在虚拟环境中运行 TensorFlow，请使用以下命令：

. tensorflow.sh

验证 TensorFlow 安装

为了验证库的功能和 GPU 支持，请在控制台中启动以下 Python 程序：

import tensorflow as tf
print(tf.reduce_sum(tf.random.normal([1000, 1000])))
print(tf.config.list_physical_devices('GPU'))

安装成功后，您将收到类似的输出，显示您的 GPU 使用情况。

(venv) user1@49069:~$ python
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
nt(tf.config.list_physical_devices('GPU'))
2024-04-22 23:39:54.472502: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
>>> print(tf.reduce_sum(tf.random.normal([1000, 1000])))
2024-04-22 23:39:55.810888: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1928] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 22282 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4090, pci bus id: 0000:07:00.0, compute capability: 8.9
tf.Tensor(332.5041, shape=(), dtype=float32)
>>> print(tf.config.list_physical_devices('GPU'))
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
>>>

此外，您可以从 此 GitHub 存储库 下载并运行用于训练神经网络的测试 Python 脚本 2c_nn_mnist_customtrain.py。

执行完成后，您应该收到类似的输出：

`bash (venv) user1@49069:~$ python ./2c_nn_mnist_customtrain.py 2024-04-23 10:35:25.832754: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. 2024-04-23 10:35:28.404381: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1928] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 22282 MB memory: -> device: 0, name: NVIDIA GeForce RTX 4090, pci bus id: 0000:07:00.0, compute capability: 8.9 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1713861338.172126 21279 service.cc:145] XLA service 0x7f36a3bd2ec0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices: I0000 00:00:1713861338.172187 21279 service.cc:153] StreamExecutor device (0): NVIDIA GeForce RTX 4090, Compute Capability 8.9 2024-04-23 10:35:38.180648: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable. 2024-04-23 10:35:38.207205: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:465] Loaded cuDNN version 8902 I0000 00:00:1713861338.355523 21279 device_compiler.h:188] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process. 1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9022 time: 14.9 test-acc: 0.932 (error: 6.75%) 1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9474 time: 4.74 test-acc: 0.953 (error: 4.71%) 1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9628 time: 4.83 test-acc: 0.963 (error: 3.67%) 1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9734 time: 4.86 test-acc: 0.968 (error: 3.25%) 1874/1875 |||||||||||||||||||||||||||||||| acc: 0.9798 time: 4.9 test-acc: 0.97 (error: 3.04%)

question_mark
Is there anything I can help you with?
question_mark
AI Assistant ×