Ububtu 安装 TensorFlow与pytorch(GPU)

Ubuntu安装TensorFlow与pytorch（GPU）

TensorFlow与pytorch介绍

TensorFlow和pytorch是两大主流的机器学习框架。其CPU版本安装较为简单。GPU版本安装问题比较多。pytorch gpu支持只需要cuda，TensorFlow gpu支持需要 cuda和cudnn

安装nvidia显卡驱动与cuda-toolkit

如果执行过apt update，请先执行apt update

sudo apt update

安装显卡驱动

# sudo apt install nvidia-driver-xxx #xxx为显卡驱动版本号，可根据当前英伟达驱动更新情况选择，如
sudo apt install nvidia-driver-510

安装cuda-toolkot

sudo apt install nvidia-cuda-toolkit

重启电脑,终端输入 nvidia-smi，有输出证明安装成功，显卡驱动会安装cuda支持，安装显卡驱动后，pytorch gpu已经得到支持。

安装TensorFlow与pytorch

对于TensorFlow，2.0版本以后，不需要单独安装TensorFlow gpu版本，直接安装TensorFlow就包含了GPU支持

#安装TensorFlow可选择pip包管理工具或conda包管理工具
pip install tensorflow #pip工具
conda install tensorflow #conda工具

对于pytorch，官网有具体的安装指令，根据cuda版本自行选择安装命令，如

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116 #cuda版本11.6

若cuda版本非11.6，请去官网自行选择

cuDNN安装

TensorFlow gpu不仅需要cuda的支持(安装显卡驱动包含了cuda驱动)，还需要cuDNN的支持,cuDNN全称 Deep Neural Network library ，深度神经网络库

安装方式1 apt安装，作者使用apt安装曾经失败过

sudo apt install nvidia-cudnn

安装方式2 源码安装
下载cudnn源码的tar包，tar包内包含两个文件家，include文件夹和lib文件夹，将两文件夹的内容分别复制到 /usr/include 和 /usr/lib

#进入解压后的cudnn文件夹
cd include
sudo cp ./* /usr/include
cd ../lib
sudo cp ./* /usr/lib

测试gpu

TensorFlow

import tensorflow as tf
tf.test.is_gpu_available()

pytorch

import torch
torch.cuda.is_available()