How To Install and Use PyTorch

Published on February 11, 2021

AI PhD Student @ UC Berkeley

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:

  1. mkdir ~/pytorch

Make a directory to hold all your assets:

  1. mkdir ~/pytorch/assets

Navigate to the pytorch directory:

  1. cd ~/pytorch

Then create a new virtual environment for the project:

  1. python3 -m venv pytorch

Activate your environment:

  1. source pytorch/bin/activate

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

  1. pip install torch torchvision

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

  1. pip install torch==1.7.1+cpu torchvision==0.8.2+cpu -f https://download.pytorch.org/whl/torch_stable.html

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:

Collecting 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:

  1. 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:

  1. 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:

  1. 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:

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

You’ll receive this output:

tensor([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:

  1. wget -O ~/pytorch/assets/imagenet_idx_to_label.json https://raw.githubusercontent.com/do-community/tricking-neural-networks/master/utils/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:

  1. wget https://raw.githubusercontent.com/do-community/tricking-neural-networks/master/step_2_pretrained.py

Note: For a more detailed walkthrough of this file step_2_pretrained.py, 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:

  1. wget -O assets/dog.jpg https://assets.digitalocean.com/articles/trick_neural_network/step2a.png

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:

  1. python step_2_pretrained.py assets/dog.jpg

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

Prediction: 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 step_2_pretrained.py 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.

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.

Thanks for learning with the DigitalOcean Community. Check out our offerings for compute, storage, networking, and managed databases.

Learn more about our products

About the authors
Default avatar
Alvin Wan


AI PhD Student @ UC Berkeley

I’m a diglot by definition, lactose intolerant by birth but an ice-cream lover at heart. Call me wabbly, witling, whatever you will, but I go by Alvin

Still looking for an answer?

Ask a questionSearch for more help

Was this helpful?
Leave a comment

This textbox defaults to using Markdown to format your answer.

You can type !ref in this text area to quickly search our full set of tutorials, documentation & marketplace offerings and insert the link!

Try DigitalOcean for free

Click below to sign up and get $200 of credit to try our products over 60 days!

Sign up

Join the Tech Talk
Success! Thank you! Please check your email for further details.

Please complete your information!

Featured on Community

Get our biweekly newsletter

Sign up for Infrastructure as a Newsletter.

Hollie's Hub for Good

Working on improving health and education, reducing inequality, and spurring economic growth? We'd like to help.

Become a contributor

Get paid to write technical tutorials and select a tech-focused charity to receive a matching donation.

Welcome to the developer cloud

DigitalOcean makes it simple to launch in the cloud and scale up as you grow — whether you're running one virtual machine or ten thousand.

Learn more
DigitalOcean Cloud Control Panel