linux系统配置nvidia显卡环境用于tensorflow、pytorch等深度学习框架

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

前言

如果机器中存在NVIDIA的GPU显卡，如何去利用呢。首先确保系统中已经安装相应型号的驱动。其次要确保python环境中有链接NVIDIA驱动相应的包，如cuda和cudnn等。如果想用NVIDIA的GPU加速，必须安装NVIDIA的驱动和cuda工具包。本文默认驱动已经安装。

一、概念

1.GPU驱动

GPU驱动是专门为图形处理单元（GPU）设计的软件，它可以让GPU能够更加高效地运行图形处理任务。GPU驱动程序通常由GPU制造商提供，例如NVIDIA和AMD。安装GPU驱动程序后，可以确保GPU能够与操作系统和其他软件兼容，并优化GPU的性能和功能。

2.cuda

CUDA（Compute Unified Device Architecture）是由NVIDIA推出的开放式并行计算平台和编程模型，它基于NVIDIA的GPU，使得开发者能够使用高级编程语言（如C/C++和CUDA C/C++）来编写程序，以利用GPU的并行处理能力。CUDA使得开发者能够更方便地利用GPU的计算能力，以加速科学计算、图形渲染、深度学习等领域的应用。

3.cuda toolkit

CUDA Toolkit，也称为NVIDIA CUDA工具包，是一个由NVIDIA开发的并行计算平台和编程模型。它提供了一套完整的开发环境，以创建经GPU加速的高性能应用。CUDA Toolkit包含多个GPU加速库、多种调试和优化工具、一个C/C++编译器以及一个用于在主要架构（包括x86、Arm和POWER）上构建和部署应用的运行时库。它支持多GPU配置中的分布式计算，科学家和研究人员可以使用它开发出可从单个GPU工作站扩展到配置数千个GPU的云端设施的应用。

4.cudnn

CuDNN是NVIDIA的深度神经网络库，它为标准例程提供了高度优化的实现，例如向前和向后卷积，池化，规范化和激活层。CuDNN使得全球的深度学习研究人员和框架开发人员能更容易地实现高性能GPU加速。CuDNN已集成到Caffe2、Chainer、Keras、MATLAB、MxNet、PyTorch和TensorFlow等广泛使用的深度学习框架中，为开发人员提供了更大的灵活性。

5.xformers

xFormers是对传统Transformer模型的一种改进和扩展。随着自然语言处理（NLP）技术的迅猛发展，Transformer模型已成为该领域的核心架构。然而，随着模型规模的不断扩大和数据量的激增，传统Transformer在计算效率和存储需求方面面临巨大挑战。为了解决这些问题，xFormers应运而生，以其高效、灵活和强大的特性受到了广泛关注。

二、安装NVIDIA显卡的cuda工具包

1.确认驱动程序的版本号

在linux命令窗口中使用nvidia-smi命令，如下图红框所示为cuda允许使用的最新版本号。

2.下载cuda工具包

打开cuda下载地址CUDA Toolkit Archive | NVIDIA Developer，选取合适的版本进行下载，如果Linux服务器处于联网的状态，可以使用wget命令进行下载，如下图所示：

wget https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run

如果Linux服务器在无联网的状态下，可以在联网的机器上直接访问如下地址进行下载，上传到Linux服务器上。

https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run

3.安装cuda工具包

使用如下命令进行安装：

sudo sh cuda_12.4.1_550.54.15_linux.run

三、安装TensorFlow、xFormers和PyTorch

1. 安装TensorFlow

• 根据cuda版本下载对应的tensorflow的GPU版本，其版本对应关系如下：

https://www.tensorflow.org/install/source_windows?hl=zh-cn#gpu

• 使用pip工具安装tensorflow，命令如下：

pip install tensorflow==xx.xx.xx

注：xx.xx.xx换成需要安装的tensorflow版本

• 安装后判断是否支持GPU

进入安装tensorflow的python环境中，输入如下命令：

import tensorflow as tf
print(len(tf.config.list_physical_devices('GPU')))

如果输出大于0，则证明安装的tensorflow是GPU版本，也同时确认安装的nvidia、cuda、cudnn驱动都是正确的。

2.安装xFormers

• 安装xFormers库依赖pytorch，根据需要查找xFormers对应的pytorch版本，如下所示：

https://github.com/facebookresearch/xformers

在main下拉框中选择对应的xformers版本，就可以得到相应版本的安装方式和对应的pytorch版本。

安装后确认是否安装成功

进入安装xformers的python环境，输入如下命令：

import xformers

如果导入xformers没有产生报错，则证明安装正确。

3.安装PyTorch

• 根据cuda版本下载对应的pytorch的GPU版本，其版本对应关系如下： 

https://pytorch.org/get-started/locally/

注：这是使用最新的pytorch

如果要使用指定版本的pytorch，如下所示：

https://pytorch.org/get-started/previous-versions/

注：使用pip安装pytorch时可能安装的是cpu版本

• 安装后判断是否支持GPU

import torch
torch.cuda.is_available()

如果输出的是True，则证明安装的pytorch是GPU版本，也同时确认安装的nvidia、cuda、cudnn驱动都是正确的。

总结

近期比较火的stable diffusion绘图工具，需要安装和配置GPU环境来提高出图的效率，本文章很好的解决了部署过程中出现的大多数问题。