How To Install and Use PyTorch

The author selected the International Medical Corps to receive a donation as part of the Write for DOnations program.


PyTorch is a framework developed by Facebook AI Research for deep learning, featuring both beginner-friendly debugging tools and a high-level of customization for advanced users, with researchers and practitioners using it across companies like Facebook and Tesla. Applications include computer vision, natural language processing, cryptography, and more. Whereas more advanced DigitalOcean articles, like Introduction to PyTorch, discuss how PyTorch works and what you can build with PyTorch, this tutorial will focus on installing PyTorch.

In this tutorial, you’ll install PyTorch’s “CPU support only” version in three steps. This installation is ideal for people looking to install and use PyTorch but who don’t have an Nvidia graphics card. In particular, you’ll install PyTorch in a Python virtual environment with virtualenv. This approach isolates the PyTorch installation, allowing you to install different PyTorch versions for each project. Once you complete the installation, you’ll validate your installation by running a short PyTorch program and then use PyTorch to perform image classification.

While you won’t need prior experience in practical deep learning or PyTorch to follow along with this tutorial, we’ll assume some familiarity with machine learning terms and concepts such as tensors. You can learn more about these concepts in An Introduction to Machine Learning.


To complete this tutorial, you will need a local development environment for Python 3 with at least 1GB of RAM. You can follow How to Install and Set Up a Local Programming Environment for Python 3 to set up Python and the essentials for your programming environment.

Step 1 — Installing PyTorch

Let’s create a workspace for this project and install the dependencies you’ll need. You’ll call your workspace pytorch:

  • mkdir ~/pytorch

Make a directory to hold all your assets:

  • mkdir ~/pytorch/assets

Navigate to the pytorch directory:

  • cd ~/pytorch

Then create a new virtual environment for the project:

  • python3 -m venv pytorch

Activate your environment:

  • source pytorch/bin/activate

Then install PyTorch. On macOS, install PyTorch with the following command:

  • pip install torch torchvision

On Linux and Windows, use the following commands for a CPU-only build:

  • pip install torch==1.7.1+cpu torchvision==0.8.2+cpu -f

Notice that you have also included torchvision by default. This sub-library includes several utilities specific to computer vision, which you’ll use later on in this tutorial. You’ll then receive output confirming PyTorch’s installation:

OutputCollecting torch   Downloading torch-1.7.1-cp38-none-macosx_10_9_x86_64.whl (108.9 MB)      |████████████████████████████████| 108.9 MB 8.3 MB/s  . . .  Successfully installed pillow-8.1.0 torch-1.7.1 torchvision-0.8.2 

Note: If you’d like to deactivate your virtual environment at any time, the command is:

  • deactivate

To reactivate the environment later, navigate to your project directory and run source pytorch/bin/activate.

Now, that you have installed PyTorch, you’ll make sure the PyTorch installation works.

Step 2 — Validating Installation

To validate the installation of PyTorch, you’re going to run a small program in PyTorch as a non-root user. Rather than creating a Python file, you’ll create this program using Python’s interactive console.

To write the program, start up your Python interpreter with the following command:

  • python3

You will receive the following prompt in your terminal:

This is the prompt for the Python interpreter, and it indicates that it’s ready for you to start entering some Python statements.

First, type this line to import the PyTorch package. Press ENTER after typing in the line of code:

  • import torch

Define a vector of zeros. For now, think of a vector as a collection of numbers, or specifically, a list of numbers. In more detail: A vector is an “arrow” in space, indicating both direction (where the arrow points), and magnitude (how long the arrow is). Next, you create a vector using a list of three numbers with Tensor(). This is a three-dimensional vector, which is an arrow in three-dimensional space.

Run the following Python:

  • torch.Tensor([0, 0, 0])

You’ll receive this output:

Outputtensor([0., 0., 0.]) 

This indicates the PyTorch installation was successful. Exit the Python interactive console by pressing CTRL+D. Next, you will build an image classifier using PyTorch.

Step 3 — Using PyTorch for Image Classification

Now that you’ve validated the PyTorch installation, you will set up an image classifier.

An image classifier accepts images as input and outputs a predicted class (like Cat or Dog). Image classifiers are the conventional “Hello World” equivalent in deep learning frameworks. Besides convention, there are a few great reasons to start with image classification. First, many image classifiers can produce predictions on a commodity CPU, without needing extensive GPU resources. Second, it is straightforward to know that your image classifier is working (or not), based on the predicted class. This is less true of other neural networks that, for example, generate text.

In this tutorial, you will use image classifiers that have already been trained. We call these pretrained image classifiers. In particular, you will use image classifiers to predict the class of an image. Prediction is alternatively called inference. In short, you will be running inference on a pretrained image classifier.

First, download a JSON file to convert neural network output to a human-readable class name:

  • wget -O ~/pytorch/assets/imagenet_idx_to_label.json

Download the following Python script, which will load an image, load a neural network with its weights, and classify the image using the neural network:

  • wget

Note: For a more detailed walkthrough of this file, please read Step 2 — Running a Pretrained Animal Classifier in the How To Trick a Neural Network tutorial.

Download the following image of a dog to run the image classifier on.

Use the following command to do so:

  • wget -O assets/dog.jpg

Finally, run the pretrained image classifier on the newly downloaded image. Pretrained means this model has already been trained and will be able to predict classes, accurately, straightaway:

  • python assets/dog.jpg

This will produce the following output, showing your animal classifier works as expected:

OutputPrediction: Pembroke, Pembroke Welsh corgi 

That concludes running inference with your pretrained model. If you’d like to use another image, you can do this by changing the first argument to your python3 command. For the argument, you’d pass in the relative path of the image file. In the next step, we’ll summarize a few recommended tools for working in PyTorch.

The PyTorch Ecosystem

In this section, we recommend several frameworks and libraries to get started with, as you work with PyTorch. Although these additional libraries may not be beneficial until later in your PyTorch journey, knowing that these tools exist is useful, so you can leverage them when you are ready.

For each deep learning library, there is a canonical higher-level framework to use. PyTorch was the exception to this rule for its first few years. However, Facebook AI Research released two frameworks in 2019 that have grown quickly in popularity. Note that neither library is fully mature as of February 2021, and as a result, native PyTorch tutorials are a much more friendly place to start.

  • Classy Vision provides utilities for starting a new project and training at scale. The library can be adapted to any application. Since the library was originally created for action recognition in videos, it contains a particularly convenient set of video utilities and out-of-the-box support for video datasets. Visit their tutorials for more. Here is an example with Classy Vision training on ImageNet: Neural-Backed Decision Trees, which you can learn more about in this tutorial on visualizing neural networks.
  • PyTorch Lightning provides a framework for creating PyTorch projects. Likewise, you can adapt this library to any application. You can use PyTorch Lightning tutorials and read their examples.

There are several deep-learning-library-agnostic tools that are additionally useful to integrate as well. Both of these are prolific in both the research and industry communities.

  • Weights and Biases is designed for experiment tracking, particularly useful for a) communicating between collaborators, b) keeping experiments organized, and c) quickly visualizing many experimental results simultaneously.
  • Tensorboard is ostensibly designed for Tensorflow. However, because PyTorch does not natively support a visualization dashboard, PyTorch practitioners adopted Tensorboard. Find the PyTorch tutorial for Tensorboard visualizations on the PyTorch website.

To deploy a PyTorch model to production, there are several commonly used options:

  • TorchScript is the natively supported deployment option for PyTorch models. However, most model deployment in 2021 uses one of the two following methods.
  • ONNX is an open-source format for all deep learning frameworks, supporting not just deployment but interoperability between deep learning frameworks. In practice, converting from a PyTorch model to the ONNX format takes some debugging to make work. This has seen incremental improvement over the years, as PyTorch increases native support for ONNX exports.
  • Caffe2 is the most commonly used deployment option of these three. Although Caffe2 was itself a standalone deep learning framework, it was subsequently merged with PyTorch in 2018. You can learn more about deploying with Caffe2 on the PyTorch tutorial website.

There are also several application-focused PyTorch GitHub repositories, useful for reproducing and then building off of state-of-the-art results in their respective areas of research:

  • Detectron2 for computer vision supports object detection, object segmentation, keypoint estimation, and more. It also includes utilities for exporting a model in either Torchscript or Caffe2.
  • Transformers for natural language processing supports a large number of state-of-the-art models and their associated tokenizers.

You can find more PyTorch libraries in the PyTorch Ecosystem.


You’ve installed PyTorch in a Python virtual environment and validated that PyTorch works by running a couple of examples. You now have the tools to explore machine learning further. For your next step in developing deep learning know-how, check out an Introduction to PyTorch: Build a Neural Network to Recognize Handwritten Digits. We’ll cover more than just PyTorch; we’ll also discuss the foundations of machine learning.

To learn more about PyTorch, see the official tutorials. You can also get started with other popular deep learning frameworks, including Tensorflow. For more machine learning content, check out our Machine Learning topic page.