How to Install PyTorch on Ubuntu 20.04 (pip & conda)

Here you will learn how to install PyTorch on Ubuntu 20.04 using either pip or conda. Specifically, you will learn how to install Python 3 and Python package manager, either pip or conda (Anaconda or Miniconda). You will also learn how to install CUDA through apt-get in the official repository of Ubuntu 20.04. Last but not least, you will be able install both the PyTorch CPU only or GPU enabled versions, even though the GPU method is strongly recommended.

PyTorch & Ubuntu 20.04

PyTorch is a popular deep learning framework written in Python. Open-sourced by Facebook, PyTorch has been used by researchers and developers for computer vision (torchvision), NLP (natural language processing, torchtext), and audio tasks.

Under the hood, PyTorch is a Tensor library (torch), similar to NumPy, and mainly provides an automatic differentiation library (torch.autograd) and a neural networks library (torch.nn). It also contains 2 components for data processing: torch.multiprocessing enables memory sharing between torch Tensors between procisesses, and torch.utils provides DataLoder class. Lastly, PyTorch also contains a component (torch.jit) to serialize and optimize models from your PyTorch code.

Code-named Focal Fossa, Ubuntu 20.04 is the most recent release of Ubuntu LTS. It is out in April 2020 and supported for 5 years until April 2025 (standard support).

One of the most exciting feature for PyTorch users is that Ubuntu 20.04 now includes the proprietary NVIDIA Linux drivers. So when you install 20.04 (or upgrade from an older Ubuntu version, such as Ubuntu 18.04), the NVIDIA Linux driver will be automatically installed.

Ubuntu 20.04 also includes Python 3 by default, so you don’t need to install it as well. This is a mandate move because the Python Foundation already announced the EOL of Python 2, which is on Jan 1, 2020. In the older Ubuntu versions, Python 2 is the default and typing python in the terminal will bring you to Python 2, which means we have to type python3 to use Python 3. Now you can install python-is-python3, which will set python to python3.

Now let’s get started. The following are the 5 steps to install PyTorch on Ubuntu 20.04. If you don’t have an NVIDIA GPU, it will take only 2 steps.

Step 1 — Install Python package manager

There are 2 major Python package managers. The first is the official one called Pip, and another one is Conda (Anaconda or Miniconda). When in doubt or for beginners, the official pip is recommended.

Install Python 3 and pip for PyTorch

Pip is the native Python package manager. We will use it to install PyTorch and torchvision. To install pip3, run the following command. Because pip3 depends on Python 3, python3 will also be installed if it is not in your Ubuntu 20.04 system.

sudo apt install python3-pip

You will see output similar to the screenshot below. Press “Y” to continue. As you may notice, it will also install python-pip-whl and python3-wheel, where wheel is the built-package format for Python.

Here is the full output for the pip3 installation:

vh@varhowto-com:~$ sudo apt install python3-pip
Reading package lists… Done
Building dependency tree
Reading state information… Done
The following additional packages will be installed:
python-pip-whl python3-wheel
The following NEW packages will be installed:
python-pip-whl python3-pip python3-wheel
0 upgraded, 3 newly installed, 0 to remove and 49 not upgraded.
Need to get 2,053 kB of archives.
After this operation, 3,455 kB of additional disk space will be used.
Do you want to continue? [Y/n] y
Get:1 http://archive.ubuntu.com/ubuntu focal/universe amd64 python-pip-whl all 20.0.2-5ubuntu1 [1,799 kB]
Get:2 http://archive.ubuntu.com/ubuntu focal/universe amd64 python3-wheel all 0.34.2-1 [23.8 kB]
Get:3 http://archive.ubuntu.com/ubuntu focal/universe amd64 python3-pip all 20.0.2-5ubuntu1 [230 kB]
Fetched 2,053 kB in 1s (2,104 kB/s)
Selecting previously unselected package python-pip-w
hl.
(Reading database … 273191 files and directories c
urrently installed.)
Preparing to unpack …/python-pip-whl_20.0.2-5ubunt
u1_all.deb …
Unpacking python-pip-whl (20.0.2-5ubuntu1) …
Selecting previously unselected package python3-whee
l.
Preparing to unpack …/python3-wheel_0.34.2-1_all.d
eb …
Unpacking python3-wheel (0.34.2-1) …
Selecting previously unselected package python3-pip.
Preparing to unpack …/python3-pip_20.0.2-5ubuntu1_
all.deb …
Unpacking python3-pip (20.0.2-5ubuntu1) …
Setting up python3-wheel (0.34.2-1) …
Setting up python-pip-whl (20.0.2-5ubuntu1) …
Setting up python3-pip (20.0.2-5ubuntu1) …
Processing triggers for man-db (2.9.1-1) …

[Alternative] Install Conda (Anoconda/Miniconda) for PyTorch

There is another popular Python package distribution called Anaconda or Miniconda. We previously written a tutorial to install Miniconda, you can read it if you prefer conda or just want to learn more. Note that conda distributes python itself, so it won’t use the system Python and you won’t need to have Python installed before installing conda.

[Alternative] Install PyTorch with CPU support only

You can skip the remaining steps and run the following 2 commands to install PyTorch and have it only use your CPU if

1. you just want to install PyTorch on your Ubuntu 20.04 machine or
2. your computer doesn’t have a GPU or
3. you only have an integrated GPU in your Intel processor (which is Intel® HD Graphics).

However, if you only use CPU, your deep learning models will run slow. Having a NVIDIA graphics card and installing PyTorch with GPU support will make your model training significantly faster.

To use pip to install PyTorch without GPU support, run

pip3 install torch==1.5.1+cpu torchvision==0.6.1+cpu -f https://download.pytorch.org/whl/torch_stable.html

To use conda (Anaconda/Miniconda) to install PyTorch without GPU support, run

conda install pytorch torchvision cpuonly -c pytorch

Step 2 — Install NVIDIA Linux driver

If you have a NVIDIA graphics card, this step will make sure you have NVIDIA’s own Linux driver installed on Ubuntu 20.04. By doing so, PyTorch can take full advantage of your GPU for processing.

First, press the windows key on you keyboard (or click Activities on the top left comer of your screen), search for “Additional Drivers”, then press enter. You will see the following window with the “Additional Drivers” tab active. Choose the entry with the biggest number. Here what we have is nvidia-driver-440. Here my graphics card is GeForce GTX 1070 (GP104), and you model will certainly vary.

You also see nvidia-driver-435 and nvidia-driver-390, don’t choose them if you are uncertain. It is worth to mention that the last entry here is “Using X.org X server — Nouveau display driver from xserver-xorg-video-nouveau (open source)“. Don’t choose this neither even though it has the open source in heart: you won’t get full advantage of your NVIDIA GPU fully — you may still try but you will feel the lag for certain.

Next, click “Apply Changes” to download and install the NVIDIA driver. Under the hood, “Software & Updates” use apt-get to achieve this. After a few minutes, it will ask you to reboot to finish the installation.

Step 3 — Install CUDA from 20.04’s official repo

PyTorch utilize CUDA for fast processing, especially autograding. As Ubuntu 20.04 has NVIDIA Linux driver built-in, we can directly install CUDA from the official repository. The current version is CUDA 10.1. To install CUDA for PyTorch on your Ubuntu 20.04 machine, run

sudo apt install nvidia-cuda-toolkit

It will download around 1.5 GB files for all the packages, and will take 3.8 GB additional disk space, so it will take a while to complete CUDA installation depending on your Internet speed. The following are a few screenshots and the complete output.

All the packages that nvidia-cuda-toolkit depends on:

Press enter to continue installation:

For your reference, here are the complete output before continuing the installation. Note that because nvidia-cuda-doc is also installed, you can view the CUDA documentation by opening /usr/share/doc/nvidia-cuda-toolkit/html/index.html in a browser.

vh@varhowto-com:~$ sudo apt install nvidia-cuda-toolkit
Reading package lists… Done
Building dependency tree
Reading state information… Done
The following additional packages will be installed:
ca-certificates-java fonts-dejavu-extra g++-8
java-common libaccinj64-10.1 libatk-wrapper-java
libatk-wrapper-java-jni libcublas10
libcublaslt10 libcudart10.1 libcufft10
libcufftw10 libcuinj64-10.1 libcupti-dev
libcupti-doc libcupti10.1 libcurand10
libcusolver10 libcusolvermg10 libcusparse10
libjs-underscore libncurses5 libnppc10
libnppial10 libnppicc10 libnppicom10
libnppidei10 libnppif10 libnppig10 libnppim10
libnppist10 libnppisu10 libnppitc10 libnpps10
libnvblas10 libnvgraph10 libnvidia-compute-435
libnvidia-ml-dev libnvjpeg10 libnvrtc10.1
libnvtoolsext1 libnvvm3 libstdc++-8-dev
libthrust-dev libtinfo5 libvdpau-dev
node-html5shiv nsight-compute nsight-systems
nvidia-cuda-dev nvidia-cuda-doc nvidia-cuda-gdb
nvidia-opencl-dev nvidia-profiler
nvidia-visual-profiler ocl-icd-opencl-dev
opencl-c-headers openjdk-8-jre
openjdk-8-jre-headless
Suggested packages:
g++-8-multilib gcc-8-doc default-jre
libstdc++-8-doc libvdpau-doc nodejs
nvidia-driver | nvidia-tesla-440-driver
| nvidia-tesla-418-driver libpoclu-dev
icedtea-8-plugin fonts-ipafont-gothic
fonts-ipafont-mincho fonts-wqy-microhei
fonts-wqy-zenhei
Recommended packages:
libnvcuvid1
The following NEW packages will be installed:
ca-certificates-java fonts-dejavu-extra g++-8
java-common libaccinj64-10.1 libatk-wrapper-java
libatk-wrapper-java-jni libcublas10
libcublaslt10 libcudart10.1 libcufft10
libcufftw10 libcuinj64-10.1 libcupti-dev
libcupti-doc libcupti10.1 libcurand10
libcusolver10 libcusolvermg10 libcusparse10
libjs-underscore libncurses5 libnppc10
libnppial10 libnppicc10 libnppicom10
libnppidei10 libnppif10 libnppig10 libnppim10
libnppist10 libnppisu10 libnppitc10 libnpps10
libnvblas10 libnvgraph10 libnvidia-compute-435
libnvidia-ml-dev libnvjpeg10 libnvrtc10.1
libnvtoolsext1 libnvvm3 libstdc++-8-dev
libthrust-dev libtinfo5 libvdpau-dev
node-html5shiv nsight-compute nsight-systems
nvidia-cuda-dev nvidia-cuda-doc nvidia-cuda-gdb
nvidia-cuda-toolkit nvidia-opencl-dev
nvidia-profiler nvidia-visual-profiler
ocl-icd-opencl-dev opencl-c-headers
openjdk-8-jre openjdk-8-jre-headless
0 upgraded, 60 newly installed, 0 to remove and 49 not upgraded.
Need to get 1,429 MB of archives.
After this operation, 3,859 MB of additional disk space will be used.
Do you want to continue? [Y/n]

Verify CUDA is installed

After the Progress is 100%, you can run nvcc -V to check if CUDA 10.1 is successfully installed on Ubuntu 20.04

nvcc -V

You will see the output in the following screenshot.

I also copied the output text for your reference:

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Sun_Jul_28_19:07:16_PDT_2019
Cuda compilation tools, release 10.1, V10.1.243

Step 4 — Install PyTorch with CUDA support

Finally, we just finished all the preparation work for installing PyTorch on Ubuntu 20.04 with NVIDIA’s CUDA GPU support.

Install PyTorch with pip

To install it using pip, run the following with the pip3 command.

pip3 install torch==1.5.1+cu101 torchvision==0.6.1+cu101 -f https://download.pytorch.org/whl/torch_stable.html

It will take quite awhile to download around 700 MB pre-built wheel files for PyTorch and its dependencies. After the progress bar of downloading is complete, you will see the following output:

Here is the complete output text:

Requirement already satisfied: future in /usr/lib/python3/dist-packages (from torch==1.5.1+cu101) (0.18.2)
Collecting numpy
Using cached numpy-1.19.0-cp38-cp38-manylinux2010_x86_64.whl (14.6 MB)
Requirement already satisfied: pillow>=4.1.1 in /usr/lib/python3/dist-packages (from torchvision==0.6.1+cu101) (7.0.0)
Installing collected packages: numpy, torch, torchvision
WARNING: The scripts f2py, f2py3 and f2py3.8 are installed in '/home/vh/.local/bin' which is not on PATH.
Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.
WARNING: The scripts convert-caffe2-to-onnx and convert-onnx-to-caffe2 are installed in '/home/vh/.local/bin' which is not on PATH.
Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.
Successfully installed numpy-1.19.0 torch-1.5.1+cu101 torchvision-0.6.1+cu101

Install PyTorch with conda

To install it using conda after you installed either Anaconda or Miniconda, run

conda install pytorch torchvision cudatoolkit=10.1 -c pytorch

Verify PyTorch Installation

Now PyTorch should be installed and have CUDA support. To double check, we will run 2 sample Python scripts. Run python3 on your Ubuntu 20.04 machine and copy & paste (or type) the code in this section.

Test PyTorch

In the following code, we create a randomly initialized tensor to test PyTorch.

import torch
print(torch.rand(5, 3))

You will see the following output printed. Yours will be similar as the ouput numbers are random.

tensor([[0.7963, 0.8564, 0.6089],
        [0.0640, 0.8827, 0.4916],
        [0.9320, 0.1722, 0.3863],
        [0.3529, 0.3285, 0.3255],
        [0.1337, 0.7180, 0.5241]])

Verify if CUDA is accessible in PyTorch

To test whether your GPU driver and CUDA are available and accessible by PyTorch, run the following Python code to determine whether or not the CUDA driver is enabled:

import torch
torch.cuda.is_available()

It should output true as seen below.

Congratulations! Now you have PyTorch with CUDA support installed on you Ubuntu 20.04 machine!

References:

• https://pytorch.org/get-started/locally/
• https://pytorch.org/get-started/previous-versions/
• https://github.com/pytorch/pytorch

4 Steps to Install PyTorch on Ubuntu 20.04

---