How to Install PyTorch on Ubuntu 22.04

Updated on November 21, 2023
How to Install PyTorch on Ubuntu 22.04 header image

Introduction

PyTorch is a machine-learning framework that includes a library of tools used to build and train deep learning models. It's highly extensible and integrates with Python to enable the fast computation of tasks. Follow the steps in this guide to install PyTorch on a Ubuntu 22.04 Server.

Prerequisites

Installation

To install PyTorch on a server, verify the system capabilities to correctly enable the framework. In addition, you can install PyTorch as a native system package or install it using an environment management tool such as Conda as described in the steps below.

Install PyTorch on a GPU Server

  1. Verify that your Server has a supported GPU driver. For example, view the Vultr NVidia GPU usage

     $ nvidia-smi

    If the above command fails, you cannot the PyTorch GPU package on the server, When successful, install the PyTorch GPU package

  2. Upgrade the Python Pip package manager

     $ pip install --upgrade pip
  3. Using Pip, install the latest PyTorch version on your server

     $ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

    The above command installs the latest PyTorch version built on the CUDA version 11.8. The additional packages, torchvision and torchaudio extend PyTorch support with image and audio processing capabilities.

Using Conda

To install PyTorch on a GPU server, either install Anaconda or Miniconda then follow the steps below.

  1. Activate your target Conda environment. For example env1

     $ conda activate env1
  2. Install the latest PyTorch version from the pytorch and the nvidia channels

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

    The above command installs the latest PyTorch version with the CUDA version 11.8. Verify the latest version and install it in your environment.

Install PyTorch on a CPU-only Server

To install PyTorch on a CPU-only server without any GPU attachment, install the latest version together with the torch, torchvision, and torchaudio processing packages as described below.

  1. Upgrade the Python Pip package manager

     $ pip install --upgrade pip
  2. Using Pip, install the latest PyTorch version

     $ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

Install with Conda

To install PyTorch using Conda on a CPU-only system, install also install the cpuonly package from the pytorch Conda channel. Because this is a CPU-only environment, do not use packages from the nvidia channel.

  1. Activate your target Conda environment. For example env1

     $ conda activate env1
  2. Install the latest PyTorch version from the pytorch channel

     $ conda install pytorch torchvision torchaudio cpuonly -c pytorch

    The above command installs PyTorch with the cpuonly, torchvision, and torchaudio packages in your Conda environment.

Test the PyTorch Installation

To verify that PyTorch is available and correctly installed on your server, perform the following test operations.

  1. Access the Python Shell

     $ python3
  2. Import the torch package

     >>> import torch
  3. Declare a random tensor

     >>> x = torch.rand(1)
  4. Print the tensor value

     >>> print(x)

    Output:

     tensor([0.4169])

    As displayed in the above output, PyTorch is actively running and performing computation tasks on your server

Test PyTorch GPU Access

  1. Access the Python Shell

     $ python3
  2. Import the PyTorch torch package

     >>> import torch
  3. Verify that PyTorch has access to the server GPU

     >>> torch.cuda.is_available()

    Output:

     True

    When the above result is True, PyTorch is correctly running with GPU access, If False, PyTorch cannot run with GPU-acceleration.

Conclusion

You have installed PyTorch on a Ubuntu server using both GPU and CPU-Only methods. Using PyTorch, you can extensively use other computation packages on your server to run and develop applications. For more information on how to install PyTorch, visit the official installation documentation.