RTX 3090,3080,3070安装 TensorFlow & PyTorch

目前（至11/7/2020）我们还无法在英伟达 RTX 30 系列 GPU（Ampere）上通过 pip/conda 安装 TensorFlow 或 PyTorch。因为这些 GPU 需要 CUDA 11.1，而当前主流的 TensorFlow/PyTorch 版本不是针对 CUDA 11.1 编译的。现在要在 30XX GPU 上运行这些库的话只能手动编译或者用英伟达 docker 容器。

另外虽然 PyTorch 支持 CUDA 11.0， 但 11.0 只支持compute capability 8.0，而不支持 RTX 30系列的compute capability 8.6。CUDA 11.1才引入了对8.6的支持(见release notes: https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html)。PyTorch + CUDA 11.0 可以在 30 系列上运行，但能用的功能集不全，性能会低不少。

解决方案

Lambda Stack，一个免费的 Ubuntu 20.04 APT 包，由 Lambda（一家设计深度学习工作站和服务器并运营公有 GPU 云的公司） 创建。Lambda Stack 可以安装并管理可在 RTX 3090，RTX 3080和 RTX 3070上运行的 TensorFlow 和 PyTorch 版本。

提醒

最新的 cuDNN 还没有针对 RTX 30 系列进行优化，一个更快的版本不久将会发布。

Lambda Stack 包括

• TensorFlow v2.3.0
• PyTorch v1.6.0
• CUDA v11.1
• cuDNN v7.6.5
• 依赖和其他框架比如 Caffe 和 Theano

系统要求

• 英伟达 GPU (如 RTX 3090, 3080, 3070, 2080 Ti 或 Quadro RTX 8000)
• Ubuntu 20.04 LTS (Focal Fossa)

安装 Lambda Stack (桌面版)

LAMBDA_REPO=$(mktemp) && \
wget -O${LAMBDA_REPO} https://lambdalabs.com/static/misc/lambda-stack-repo.deb && \
sudo dpkg -i ${LAMBDA_REPO} && rm -f ${LAMBDA_REPO} && \
sudo apt-get update && sudo apt-get install -y lambda-stack-cuda

然后重启。

安装 Lambda Stack (服务器版)

LAMBDA_REPO=$(mktemp) && \
wget -O${LAMBDA_REPO} https://lambdalabs.com/static/misc/lambda-stack-repo.deb && \
sudo dpkg -i ${LAMBDA_REPO} && rm -f ${LAMBDA_REPO} && \
sudo apt-get update && \
sudo apt-get --yes upgrade && \
sudo apt-get install --yes --no-install-recommends lambda-server && \
sudo apt-get install --yes --no-install-recommends nvidia-headless-455 && \
sudo apt-get install --yes --no-install-recommends lambda-stack-cuda

然后重启。

开始使用 TensorFlow/PyTorch

$ python
>>> import tensorflow
>>> tensorflow.__version__
'2.3.0'
>>> import torch
>>> torch.__version__
'1.6.0'

管理你的 TensorFlow/PyTorch 安装

如果有更新版本的 PyTorch、 TensorFlow、 CUDA 或 cuDNN 发布，只需运行:

sudo apt-get update && sudo apt-get dist-upgrade

然后重启。

参考：https://lambdalabs.com/blog/install-tensorflow-and-pytorch-on-rtx-30-series/