Linux Ubuntu配置GPU，并运行tensorflow demo

little_boy_lm

于 2024-07-04 17:06:46 发布

阅读量317

点赞数 4

分类专栏：踩坑日志文章标签： linux ubuntu tensorflow

本文链接：https://blog.csdn.net/little_boy_lm/article/details/140182265

版权

踩坑日志专栏收录该内容

4 篇文章 0 订阅

订阅专栏

一、运行环境

Ubuntu 22.04.4 LTS tensorflow-2.16.1 CUDA-12.3 cuDNN-8.9

确定版本，下图是tensorflow测试过的版本，按需选择版本

二、安装CUDA

至官网，下载合适的版本

wget https://developer.download.nvidia.com/compute/cuda/12.3.1/local_installers/cuda_12.3.1_545.23.08_linux.run
sudo sh cuda_12.3.1_545.23.08_linux.run

三、cuDNN的安装

下载cuDNN

官网链接，注意这里需要登录nvidia才能下载

解压

#解压
xz -d -c cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz | tar xf -
#拷贝文件
cd cudnn-linux-x86_64-8.9.7.29_cuda11-archive
sudo cp include/cudnn*    /usr/local/cuda/include 
sudo cp lib/libcudnn*    /usr/local/cuda/lib64 
#修改权限
sudo chmod a+r /usr/local/cuda/include/cudnn*   /usr/local/cuda/lib64/libcudnn*

四、tensorflow

安装

python3 -m pip install tensorflow[and-cuda] -i https://pypi.tuna.tsinghua.edu.cn/simple

测试

python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

最后一行如果显示如下信息，那么恭喜你安装成功！

简单demo

import tensorflow as tf

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

# 确保TensorFlow使用GPU
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 设置TensorFlow增长式分配GPU内存
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

# 创建两个形状为[2, 3]和[3, 2]的矩阵
a = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32)
b = tf.constant([[1, 2], [3, 4], [5, 6]], dtype=tf.float32)

# 进行矩阵乘法
c = tf.matmul(a, b)

print(c)

参考

Linux安装cuda和cudnn教程-CSDN博客

https://www.tensorflow.org/install

【Linux】监控NVIDIA GPU显卡占用状态的命令_显卡消耗检测 linux-CSDN博客