目前（至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/