This post shows you how to install TensorFlow & PyTorch (and all dependencies) in under 2 minutes using Lambda Stack, a freely available Ubuntu 20.04 APT package created by Lambda (we design deep learning workstations & servers and run a public GPU Cloud)

What will be installed

  • TensorFlow v2.4.1
  • PyTorch v1.8.0
  • CUDA v11.1
  • cuDNN v8.0.5
  • Dependencies and other frameworks like Caffe and Theano

System requirements

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

Install Lambda Stack for a 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-cuda

Reboot your machine.

Install Lambda Stack for a server

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-460 && \
sudo apt-get install --yes --no-install-recommends lambda-stack-cuda

Reboot your machine.

Start using TensorFlow/PyTorch

$ python
>>> import tensorflow
>>> tensorflow.__version__
'2.5.0'
>>> import torch
>>> torch.__version__
'1.8.1'

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?