tensorflow2 安装 | ubuntu 20.04, cuda 12.2, RTX 4090

各种版本：

OS: Ubuntu 20.04

GPU: RTX4090

Cuda: 12.2

Driver Version: 535.183.01

1. 检查对应的 tensorflow version

Tensorflow与Python、CUDA、cuDNN的版本对应表

选择：tensorflow-2.15.0, python 3.9, cudnn 8.9.

tf 2.13.0及以后的版本都不需要额外安装tensorflow-gpu

2. 安装tf==2.15.0

conda create -n tf2_env python=3.9

conda activate tf2_env

pip install tensorflow==2.15.0

python -c "import tensorflow as tf; print(tf.__version__)"

3.安装cudnn==8.9

下载cudnn；

Download cuDNN v8.9.7 (December 5th, 2023), for CUDA 12.x

解压下载的cudnn-linux-x86_64-8.9.7.29_cuda12-archive之后

sudo cp cudnn-linux-x86_64-8.9.7.29_cuda12-archive/include/cudnn*.h /usr/local/cuda/include
sudo cp cudnn-linux-x86_64-8.9.7.29_cuda12-archive/lib/libcudnn* /usr/local/cuda/lib64

检查cudnn的安装

cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

输出

cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
#define CUDNN_MAJOR 8
#define CUDNN_MINOR 9
#define CUDNN_PATCHLEVEL 7
--
#define CUDNN_VERSION (CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)

/* cannot use constexpr here since this is a C-only file */

4. 检查tensorflow识别gpu

import tensorflow as tf

# 打印可用的物理设备列表
print("Available physical devices:")
print(tf.config.list_physical_devices())

# 检查是否可以找到 GPU
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    print(f"GPU(s) found: {len(gpus)}")
    for gpu in gpus:
        print(gpu)
else:
    print("No GPU found")

输出

python test.py 
2024-09-02 13:28:16.796042: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable TF_ENABLE_ONEDNN_OPTS=0.
2024-09-02 13:28:16.818693: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-09-02 13:28:16.818723: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-09-02 13:28:16.819410: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-09-02 13:28:16.823037: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-09-02 13:28:17.202044: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Available physical devices:
2024-09-02 13:28:17.472883: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:901] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-09-02 13:28:17.524635: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:901] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-09-02 13:28:17.524748: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:901] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
GPU(s) found: 1
PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')