Installing Pytorch with GPU Support (CUDA) in Ubuntu 18.04 — Complete Guide

PyTorch (Image Source — https://pytorch.org/get-started/locally/)

These days most of the research level machine learning algorithms are coded to be run on CUDA enabled GPUs due to the clear advantage at processing the networks at greater speeds, especially when it comes to ‘Computer Vision’ problems. Before moving into coding and running the benchmarks using PyTorch, we need to setup the environment to use the GPU in processing our networks.

PyTorch is a more flexible framework than TensorFlow which can be used in such applications and it’s widely used in Research level applications than TensorFlow. Therefore, setting up an environment for PyTorch which is capable of using the GPU for its computation is important.

Check my TensorFlow GPU installation in Ubuntu 18.04 if you want to test that out as well (Click Here)

This process is a lot more straightforward than for TensorFlow and lets take a look at how it can be done in Ubuntu 18.04.

Important:

1. I have demonstrated this using “Anaconda” and I have installed the packages by creating a separate environment for PyTorch. This can come in handy if you ever feel like removing the PyTorch installation or if you run into errors which seem unsolvable, you can just remove the environment and start again by creating a new environment without removing the whole anaconda installation altogether.
2. I have mentioned the possible errors that can appear in your terminal when you run the following commands and how to correct them alongside each of those commands. If you run into errors that are not addressed in this article, please comment them below and I will try to help you to the best I can. Stack-overflow can also be really useful in such situations.
3. At last, I have given you example codes to run which you can use to check whether the installation was successful. Trying them in a ‘jupyter notebook’ is recommended.

Step 01 : System Check

(i) Make sure your system is CUDA capable

Visit the following sites and make sure your GPU is CUDA capable. https://developer.nvidia.com/cuda-gpus

But this list of CUDA Enabled GPUs seems to be incomplete because the GPU that I am using right now is CUDA enabled and it is not listed in the list given above for some reason!

Therefore, please check thoroughly before moving ahead with this process.

Googling “Is ___ CUDA enabled?” with your GPU version in the blanks might give you the answer you want.

Or else, visit the official promotional page of your GPU by NVIDIA and see whether it is listed as CUDA enabled.

(ii) Identify which CUDA/ cuDNN packages are compatible with your GPU

Different GPUs support different versions of CUDA and if your GPU is relatively old, then there’s a higher chance that it might not work with the latest versions of CUDA. Therefore, checking the appropriate versions which works with your own GPU is important before installing the latest CUDA Toolkit version blindly. This can be done searching on the internet or visiting the official NVIDIA developer forum.

Step 02 : Install PyTorch

(i) Install Anaconda

Visit https://www.anaconda.com/products/individual and install the version of your preference.

(ii) Create a New Environment for Pytorch

Please note that I have used python version 3.7 for this environment based on compatibility issues. Replace that with the version that your machine is compatible for a clean installation. Please not that you need to have python 3.7 or higher to work with PyTorch.

conda create -n torch-gpu python=3.7
source activate torch-gpu

— — — — — — — — Possible Error at this point: — — — — — — — —

conda: command not found

This occurs because the path for anaconda installation has not been set in your .bashrc or .zshrc

Try:

export PATH="/home/username/anaconda3/bin:$PATH"

— — — — — — — — -End of the Error Resolving — — — — — — — —

(iii) Install PyTorch

For PyTorch it is straight forward than TensorFlow installation because you don’t have to separately install CUDA ToolKit and cuDNN because you can install them at once using a single command as given in PyTorch official website.

You only need to identify the appropriate CUDA ToolKit version and the cuDNN package will be automatically selected unlike in TensorFlow installation which makes the PyTorch installation very easy.

Visit https://pytorch.org/get-started/locally/ and the site will automatically select the appropriate pytorch version depending on your operating system, its version and CUDA version. If it needs to be changed, you can always change the selected tiles and make it right.

Source — https://pytorch.org/get-started/locally/

If your system is relatively old, you can check the Previous Versions of PyTorch and choose which is suitable for your system. Run the command given by the PyTorch website inside the already activated environment which we created for PyTorch.

Example Code:

conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch

Step 03 : Validate the Installation

Run the following the following in a jupyter notebook validatethe installation.

import torch
torch.cuda.is_available()

or

import torch
torch.cuda.device(0)

or

import torch
torch.cuda.device_count()

or

import torch
torch.cuda.get_device_name(0)

Enjoy !

:)

Check my TensorFlow GPU installation on Ubuntu 18.04 if you want to test that out as well (Click Here)