Currently (as of 11/7/2020) we cannot install TensorFlow or PyTorch on Nvidia RTX 30 Series Gpus (Ampere) via PIP/Conda. Because these Gpus require CUDA 11.1, the current mainstream TensorFlow/PyTorch versions are not compiled for CUDA 11.1. These libraries can now be run on a 30XX GPU only by manual compilation or using a Nvidia Docker container.

The solution

Lambda Stack, a free Ubuntu 20.04 APT package, was created by Lambda, a company that designs deep learning workstations and servers and operates the public GPU cloud. Lambda Stack installs and manages versions of TensorFlow and PyTorch that run on the RTX 3090, RTX 3080, and RTX 3070.

remind

The latest cuDNN has not been optimized for the RTX 30 series, and a faster version will be released soon.

Lambda Stack including

  • TensorFlow v2.3.0
  • PyTorch v1.6.0
  • CUDA v11.1
  • CuDNN v7.6.5
  • Dependencies and other frameworks such as 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)

Installing Lambda Stack (Desktop)

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-cudaCopy the code

Then restart.

Installing Lambda Stack (Server Edition)

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-cudaCopy the code

Then restart.

Start using TensorFlow/PyTorch

$python >>> import tensorFlow >>> tensorflow.__version__ '2.3.0' >> Import torch >> torch.Copy the code

Manage your TensorFlow/PyTorch installation

If an updated version of PyTorch, TensorFlow, CUDA, or cuDNN is available, simply run:

sudo apt-get update && sudo apt-get dist-upgrade
Copy the code

Then restart.

Reference:

Lambdalabs.com/blog/instal…