How to Install PyTorch on Ubuntu

7 minutes

Dwijadas Dey

Introduction

PyTorch is a powerful and open-source deep learning platform created by Facebook's AI research division, using Python and the Torch library. It's a popular choice among research and academic teams for developing and training deep learning models in artificial intelligence (AI), machine learning (ML), and data science. PyTorch offers GPU support, which speeds up the process of inference and training for machine learning models.

The versatility of PyTorch is derived from its ability to generate a dynamic computational graph - a feature to generate a graph on the fly during runtime. This makes PyTorch a better option for generating flexible and intuitive models as compared to static computation graphs.

This article will guide you on installing PyTorch on Ubuntu 22.04 via PIP and Conda. The installation procedure will also cover how to use Conda to install GPU enabled PyTorch in Ubuntu 22.04. Towards the end of this article we will also touch upon the advantages of using PyToch.

Prerequisites

To install PyTorch on Ubuntu 22.04, ensure you meet the following requirements.

A virtual machine running Ubuntu 22.04.
You have SSH access to the virtual machine with sudo privilege.
The Latest version of PyTorch requires Python 3.8 or later.

PyTorch - GPU or CPU version?

The PyTorch has two versions - a CPU compiled version and a GPU compiled version which means PyTorch is hardware agnostic and can operate on various hardware platforms. Therefore depending on your computing and hardware requirements, your experience with PyTorch will vary in terms of processing time.

To avail the full benefits of PyTorchs CUDA or ROCm, it is required to have NVIDIA or AMD GPU support in your server which can speed up both training and inference from models tenfold.

However, if GPU support is not available in your system, you can install the CPU compiled PyTorch. In that case, You will miss out on the benefits of GPU acceleration but still can use PyTorch on your system.

Install CPU compiled PyTorch

Installing PyTorch via Conda

Installing PyTorch through Conda necessitates that Conda is present on your system. However, if it is not available on your system, install Conda by downloading and running the installer from its official site using the steps given in the next section.

$ cd /home/ubuntu
$ wget https://repo.anaconda.com/archive/Anaconda3-2024.06-1-Linux-x86_64.sh
$ sudo chmod u+x Anaconda3-2024.06-1-Linux-x86_64.sh
$ sudo ./Anaconda3-2024.06-1-Linux-x86_64.sh
...
...
Please answer 'yes' or 'no':'
>>> yes
Anaconda3 will now be installed into this location:
/root/anaconda3
 - Press ENTER to confirm the location
 - Press CTRL-C to abort the installation
 - Or specify a different location below
[/root/anaconda3] >>> /home/ubuntu/Anaconda
PREFIX=/home/ubuntu/Anaconda
Unpacking payload ...
...
...

Review the license and type 'yes' to proceed with the installation of Anaconda. Update the prefix path for Anaconda installation if you want to install it in a different location. Here we are using the location /home/ubuntu/Anaconda as a prefix path for Anaconda installation.

Once the installation is over, update the PATH environment variable with the path of the Anaconda binary.

$ echo 'export PATH="/home/ubuntu/Anaconda/bin:$PATH"' >> .bashrc
$ source .bashrc
$ conda --version
conda 24.5.0
$ conda init  ← Initialize Anaconda.

Next, create and activate a virtual environment in the system but make sure you have already executed 'conda init' in the terminal.

$ conda create -n  env1
$ conda activate env1
(env1)$

Now install the latest PyTorch version with CPU support only by utilizing the PyTorch channel.

$ conda install pytorch torchvision torchaudio cpuonly -c pytorch
Channels:
- pytorch
- defaults
Platform: linux-64
Collecting package metadata (repodata.json): done
Solving environment: done
## Package Plan ##
...
...

Installing PyTorch via PIP

If you don't have access to Anaconda then it is also possible to install PyTorch using Python pip. Therefore, before installing PyTorch using pip, ensure Python PIP is installed on your system or, if it's already present, update the Pip package.

$ sudo apt install python3-pip
OR
$ sudo pip3 install --upgrade pip
[sudo] password for ubuntu: 
Requirement already satisfied: pip in /usr/local/lib/python3.10/dist-packages (24.2)

Now install PyTorch with CPU support only.

$ conda install pytorch torchvision torchaudio cpuonly -c pytorch

Test PyTorch Installation

To test PyTorch installation has gone correctly on your Ubuntu server, print a random tensor value by importing the PyTorch module in the Python shell.

(env1) ubuntu@ip-172-26-4-236:~$ python3
Python 3.12.4 | packaged by Anaconda, Inc. | (main, Jun 18 2024, 15:12:24) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> x = torch.rand(1)
>>> print(x)
tensor([0.3501])
>>>

The output of the above print statement confirms that the PyTorch is running and performing the computational task correctly.

To check the installed PyTorch's version in your Ubuntu server, run the following command.

$ python3
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch;
>>> 
>>> torch.__version__
'2.4.0+cpu'
>>>

Install PyTorch on a GPU Server

While installing GPU compiled PyTorch on an Ubuntu 22.04 server, make sure your server has the supported hardware and a corresponding GPU driver. For example, if your server is using an NVIDIA graphics card then verify the GPU setting by viewing the statistics about GPU performance using either of the following commands.

$ nvidia-smi
or
$ watch -n 1 nvidia-smi

The above command will tell you the driver API version and works just fine in most cases. However, if the above command fails then you have to troubleshoot and install NVIDIA GPU driver on your Ubuntu server.

Installing PyTorch via Conda

To install PyTorch with CUDA using Conda, you need to have Conda on your server. In case, Conda is not available in your system - install it by following the steps from here.

Once Conda is installed, activate the target environment in the server.

$ conda activate env1

Install the latest PyTorch with CUDA version 11.8 with Nvidia GPU support by utilizing the PyTorch and the Nvidia channels. The CUDA runtime libraries, cuDNN, NCCL, etc. will be installed along with the PyTorch binary.

$ conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

If the command above does not matches with your environment, generate the command from the PyTorch site. On successful execution of the above command, the latest PyTorch with CUDA version 11.8 will be installed on your Ubuntu server.

Installing PyTorch via PIP

To install PyTorch using Python PIP, update the Python package manager(pip) and then install the latest PyTorch with CUDA version 11.8 using the following command.

$ sudo apt install python3-pip
OR
$ sudo pip3 install --upgrade pip
$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

The above Python pip command will install PyTorch with CUDA version 11.8 on your Ubuntu server.

Test PyTorch Installation

Firstly, import the torch package to test if PyTorch installation is correct and accessible, and then test if the GPU is accessible from PyTorch using PyTorch's generic method torch.cuda. This method will return either True or False depending on whether CUDA is accessible or not.

$ python3 -c 'import torch; print(torch.cuda.is_available())'
True

Uninstalling PyTorch

To uninstall PyTorch which has been installed using Conda, run the following command from the Python virtual environment.

$ conda remove pytorch -y

Use the following pip command to remove PyTorch, if you've set up the same through PIP.

$ pip3 uninstall torch -y

Advantages of PyTorch

Dynamic Computational Graphs

A unique feature of PyTorch that allows on-the-fly model definition and modification. This flexible feature empowers developers and researchers to explore complex deep-learning models by adjusting the computational graph while the model is in operation.

Ease of Use

PyTorch’s user-friendly API and the dynamic computational graphs are easy to learn like standard Python code. The simplicity and ease of use in developing deep learning models using PyTorch make it an attractive choice for not only beginners but also for experienced professionals.

GPU acceleration

PyTorch’s native support for CUDA allows users to make full use of GPU acceleration for training and inference of models and speeds up complex deep learning tasks.

Industry Adoption

PyTorch’s performance, ease of use, and ever-growing ecosystem have catapulted its adoption in both academics and industries. Research institutes and leading tech companies are using it in areas such as computer vision, NLP, and more.

Seamless integration with Other Tools

PyTorch integrates easily with other popular Python libraries such as SciPy, NumPy, and pandas. Moreover, PyTorch is interoperable with other frameworks and libraries through projects like ONNX(Open Neural Network Exchange).

Conclusion

At this point, You have successfully installed PyTorch on Ubuntu 22.04 using both Conda and PIP. Moreover, this guide also shows you how to install a GPU compiled PyTorch depending on the hardware platform you are using.

With PyTorch up and running, you now have access to a robust machine-learning framework for developing and training neural networks, thereby unlocking a plethora of opportunities for research and trial-and-error in the field of machine learning and data science.

How to install and configure Wireguard VPN Server on Almalinux 9 and Rocky Linux 9

Tips and Tricks to pass the CKA and CKAD examination

Comments