![7fcd31af25d8d50c23e498111356b58f.png](https://i-blog.csdnimg.cn/blog_migrate/6a22e99a5b2c778e5e984ba5a89a2e4c.png)
目前(至11/7/2020)我们还无法在英伟达 RTX 30 系列 GPU(Ampere)上通过 pip/conda 安装 TensorFlow 或 PyTorch。因为这些 GPU 需要 CUDA 11.1,而当前主流的 TensorFlow/PyTorch 版本不是针对 CUDA 11.1 编译的。现在要在 30XX GPU 上运行这些库的话只能手动编译或者用英伟达 docker 容器。
解决方案
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
系统要求
- NVIDIA 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/