温馨提示×

如何在 Ubuntu 上安装 TensorFlow

小樊
37
2025-12-24 09:48:53
栏目: 智能运维

在 Ubuntu 上安装 TensorFlow 的完整指南

一 环境准备与版本选择

  • 建议使用 Ubuntu 20.04/22.04 LTS,并优先在 Python 3.7–3.11 的虚拟环境中安装,避免依赖冲突。
  • 硬件与驱动:如需 GPU 加速,需 NVIDIA 显卡且驱动支持目标 CUDA 版本;可用 nvidia-smi 查看驱动与最高支持的 CUDA 版本。
  • 版本匹配要点(示例):
    • TensorFlow 2.10–2.12CUDA 11.8cuDNN 8.6
    • TensorFlow 2.13+CUDA 12.1cuDNN 8.9
  • 安装方式概览:原生 pipDocker 镜像、或 Anaconda/venv 虚拟环境。以上要点与匹配关系可显著提升安装成功率与稳定性。

二 原生 pip 安装步骤(CPU 或 GPU)

  • 1)创建并激活虚拟环境(推荐)
    • 使用 venv:
      • python3 -m venv tf-env && source tf-env/bin/activate
    • 使用 conda:
      • conda create -n tf-env python=3.9 -y && conda activate tf-env
    • 升级 pip:pip install -U pip
  • 2)安装 TensorFlow
    • CPU 与 GPU 的统一包(推荐):pip install tensorflow(自 2.x 起该包同时包含 CPU 与 GPU 支持,实际是否启用 GPU 取决于系统是否配置好 CUDA/cuDNN 与驱动)。
    • 如确需单独安装 GPU 包(部分教程仍使用):pip install tensorflow-gpu,但务必与 CUDA/cuDNN 版本匹配。
  • 3)验证安装
    • python - <<'PY' import tensorflow as tf print("TF version:", tf.__version__) print("GPUs:", tf.config.list_physical_devices("GPU")) PY
      若输出包含 GPU 设备信息,则 GPU 可用。以上流程与命令适用于 Ubuntu 上的 pip 安装场景。

三 GPU 支持与 Docker 方式(更省心)

  • 使用 Docker 是在 GPU 上运行 TensorFlow 的最简方式之一:主机仅需安装 NVIDIA 驱动,无需在主机安装完整的 CUDA Toolkit
  • 步骤:
    • 安装 NVIDIA Container Toolkit(添加 GPG 与 APT 源后安装 nvidia-container-toolkit)。
    • 验证 GPU 与容器运行时:
      • docker run --gpus all --rm nvidia/cuda nvidia-smi
    • 拉取并运行带 GPU 的 TensorFlow 镜像(含 Jupyter):
      • docker pull tensorflow/tensorflow:latest-gpu-jupyter
      • docker run -it -p 8888:8888 tensorflow/tensorflow:latest-gpu-jupyter
    • 浏览器访问:http://localhost:8888。如需长期运行与自动重启,可加上 --restart=always --name tf-gpu 等参数。该方式隔离性好、依赖冲突少,适合快速上手与复现实验。

四 常见问题与排查

  • 驱动与兼容性:用 nvidia-smi 确认驱动正常且支持的 CUDA 版本;TensorFlow 与 CUDA/cuDNN 版本必须匹配(如 2.12→CUDA 11.8/cuDNN 8.62.13+→CUDA 12.1/cuDNN 8.9)。
  • 库路径问题:若报找不到 CUDA 动态库(如 libcudart.so),检查并设置环境变量:
    • export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
    • export CUDA_HOME=/usr/local/cuda
  • GPU 不可见:在 Python 中确认 tf.config.list_physical_devices('GPU') 是否返回设备;若为空,优先检查驱动、CUDA/cuDNN 版本匹配与容器运行时配置。
  • 显存不足:训练时报 OOM,可减小 batch_size、启用内存按需增长:
    • gpus = tf.config.experimental.list_physical_devices('GPU')
    • tf.config.experimental.set_memory_growth(gpus[0], True)
  • 虚拟环境未生效:确保已激活对应的 venv/conda 环境后再 pip install 与运行 python。以上排查思路覆盖最常见的驱动、版本与资源问题。

0