Install TensorFlow & PyTorch for the RTX 3090, 3080, 3070

As of 11/6/2020, you can't pip/conda install a TensorFlow or PyTorch version that runs on NVIDIA's RTX 30 series GPUs (Ampere). These GPUs require CUDA 11.1, and the current TensorFlow/PyTorch releases aren't built against CUDA 11.1. Right now, getting these libraries to work with 30XX GPUs requires manual compilation or NVIDIA docker containers.

Solution
Lambda Stack, a freely available Ubuntu 20.04 APT package created by Lambda (we design deep learning workstations and servers and run a public GPU Cloud). Lambda Stack installs and manages versions of TensorFlow and PyTorch than can run on the RTX 3090, RTX 3080, and RTX 3070.

Caveat
The latest cuDNN isn't optimized for the RTX 30 series. A faster version will be released in the near future.

Lambda Stack includes

  • TensorFlow v2.3.0
  • PyTorch v1.6.0
  • CUDA v11.1
  • cuDNN v7.6.5
  • Dependencies and other frameworks like Caffe and Theano

System requirements

  • An NVIDIA GPU (e.g. RTX 3090, 3080, 3070, 2080 Ti, Quadro RTX 8000)
  • Ubuntu 20.04 LTS (Focal Fossa)

Install Lambda Stack (desktop version)

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

Reboot your machine.

Install Lambda Stack (server version)

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

Reboot your machine.

Start using TensorFlow/PyTorch

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

Manage your TensorFlow/PyTorch installation

If a new version of PyTorch, TensorFlow, CUDA, cuDNN is released, simply run:

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

Reboot your machine.

Interested in a GPU workstation, server, or cloud instance?

!-- Intercom -->