PyTorch vs TensorFlow for Your Python Deep Learning Project

Who Uses TensorFlow?

TensorFlow has a reputation for being a production-grade deep learning library. It has a large and active user base and a proliferation of official and third-party tools and platforms for training, deploying, and serving models.

After PyTorch was released in 2016, TensorFlow declined in popularity. But in late 2019, Google released TensorFlow 2.0, a major update that simplified the library and made it more user-friendly, leading to renewed interest among the machine learning community.

Code Style and Function

Before TensorFlow 2.0, TensorFlow required you to manually stitch together an abstract syntax tree—the graph—by making tf.* API calls. It then required you to manually compile the model by passing a set of output tensors and input tensors to a session.run() call.

A Session object is a class for running TensorFlow operations. It contains the environment in which Tensor objects are evaluated and Operation objects are executed, and it can own resources like tf.Variable objects. The most common way to use a Session is as a context manager.

In TensorFlow 2.0, you can still build models this way, but it’s easier to use eager execution, which is the way Python normally works. Eager execution evaluates operations immediately, so you can write your code using Python control flow rather than graph control flow.

To see the difference, let’s look at how you might multiply two tensors using each method. Here’s an example using the old TensorFlow 1.0 method:

>>> import tensorflow as tf

>>> tf.compat.v1.disable_eager_execution()

>>> x = tf.compat.v1.placeholder(tf.float32, name = "x")
>>> y = tf.compat.v1.placeholder(tf.float32, name = "y")

>>> multiply = tf.multiply(x, y)

>>> with tf.compat.v1.Session() as session:
...     m = session.run(
...         multiply, feed_dict={x: [[2., 4., 6.]], y: [[1.], [3.], [5.]]}
...     )
...     print(m)
[[ 2.  4.  6.]
 [ 6. 12. 18.]
 [10. 20. 30.]]

This code uses TensorFlow 2.x’s tf.compat API to access TensorFlow 1.x methods and disable eager execution.

You first declare the input tensors x and y using tf.compat.v1.placeholder tensor objects. Then you define the operation to perform on them. Next, using the tf.Session object as a context manager, you create a container to encapsulate the runtime environment and do the multiplication by feeding real values into the placeholders with a feed_dict. Finally, still inside the session, you print() the result.

With eager execution in TensorFlow 2.0, all you need is tf.multiply() to achieve the same result:

>>> import tensorflow as tf

>>> x = [[2., 4., 6.]]
>>> y = [[1.], [3.], [5.]]
>>> m = tf.multiply(x, y)

>>> m
<tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[ 2.,  4.,  6.],
       [ 6., 12., 18.],
       [10., 20., 30.]], dtype=float32)>

In this code, you declare your tensors using Python list notation, and tf.multiply() executes the element-wise multiplication immediately when you call it.

If you don’t want or need to build low-level components, then the recommended way to use TensorFlow is Keras. It has simpler APIs, rolls common use cases into prefabricated components for you, and provides better error messages than base TensorFlow.

Special Features

TensorFlow has a large and well-established user base and a plethora of tools to help productionize machine learning. For mobile development, it has APIs for JavaScript and Swift, and TensorFlow Lite lets you compress and optimize models for Internet of Things devices.

You can get started using TensorFlow quickly because of the wealth of data, pretrained models, and Google Colab notebooks that both Google and third parties provide.

Many popular machine learning algorithms and datasets are built into TensorFlow and are ready to use. In addition to the built-in datasets, you can access Google Research datasets or use Google’s Dataset Search to find even more.

Keras makes it easier to get models up and running, so you can try out new techniques in less time. Indeed, Keras is the most-used deep learning framework among the top five winningest teams on Kaggle.

One drawback is that the update from TensorFlow 1.x to TensorFlow 2.0 changed so many features that you might find yourself confused. Upgrading code is tedious and error-prone. Many resources, like tutorials, might contain outdated advice.

PyTorch doesn’t have the same large backward-compatibility problem, which might be a reason to choose it over TensorFlow.

Tensorflow Ecosystem

Some highlights of the APIs, extensions, and useful tools of the TensorFlow extended ecosystem include:

• TensorFlow Hub, a library for reusable machine learning modules
• Model Garden, an official collection of models that use TensorFlow’s high-level APIs
• Hands-On Machine Learning with Scikit-Learn, Keras, and TensorFlow, a comprehensive introduction to machine learning using TensorFlow

Who Uses PyTorch?

PyTorch has a reputation for being more widely used in research than in production. However, since its release the year after TensorFlow, PyTorch has seen a sharp increase in usage by professional developers.

The 2020 Stack Overflow Developer Survey list of most popular “Other Frameworks, Libraries, and Tools” reports that 10.4 percent of professional developers choose TensorFlow and 4.1 percent choose PyTorch. In 2018, the percentages were 7.6 percent for TensorFlow and just 1.6 percent for PyTorch.

As for research, PyTorch is a popular choice, and computer science programs like Stanford’s now use it to teach deep learning.

Code Style and Function

PyTorch is based on Torch, a framework for doing fast computation that is written in C. Torch has a Lua wrapper for constructing models.

PyTorch wraps the same C back end in a Python interface. But it’s more than just a wrapper. Developers built it from the ground up to make models easy to write for Python programmers. The underlying, low-level C and C++ code is optimized for running Python code. Because of this tight integration, you get:

• Better memory and optimization
• More sensible error messages
• Finer-grained control of model structure
• More transparent model behavior
• Better compatibility with NumPy

That means you can write highly customized neural network components directly in Python without having to use a lot of low-level functions.

PyTorch’s eager execution, which evaluates tensor operations immediately and dynamically, inspired TensorFlow 2.0, so the APIs for both look a lot alike.

Converting NumPy objects to tensors is baked into PyTorch’s core data structures. That means you can easily switch back and forth between torch.Tensor objects and numpy.array objects.

For example, you can use PyTorch’s native support for converting NumPy arrays to tensors to create two numpy.array objects, turn each into a torch.Tensor object using torch.from_numpy(), and then take their element-wise product:

>>> import torch
>>> import numpy as np

>>> x = np.array([[2., 4., 6.]])
>>> y = np.array([[1.], [3.], [5.]])

>>> m = torch.mul(torch.from_numpy(x), torch.from_numpy(y))

>>> m.numpy()
array([[ 2.,  4.,  6.],
       [ 6., 12., 18.],
       [10., 20., 30.]])

Using torch.Tensor.numpy() lets you print out the result of matrix multiplication—which is a torch.Tensor object—as a numpy.array object.

The most important difference between a torch.Tensor object and a numpy.array object is that the torch.Tensor class has different methods and attributes, such as backward(), which computes the gradient, and CUDA compatibility.

Special Features

PyTorch adds a C++ module for autodifferentiation to the Torch backend. Autodifferentiation automatically calculates the gradient of the functions defined in torch.nn during backpropagation.

By default, PyTorch uses eager mode computation. You can run a neural net as you build it, line by line, which makes it easier to debug. It also makes it possible to construct neural nets with conditional execution. This dynamic execution is more intuitive for most Python programmers.

PyTorch Ecosystem

Some highlights of the APIs, extensions, and useful tools of the PyTorch extended ecosystem include:

• The fast.ai API, which makes it extremely easy to build models quickly
• TorchServe, an open source model server developed in collaboration between AWS and Facebook
• TorchElastic for training deep neural networks at scale using Kubernetes
• PyTorch Hub, an active community for sharing and extending cutting-edge models

Style

If you’re a Python programmer, then PyTorch will feel easy to pick up. It works the way you’d expect it to, right out of the box.

On the other hand, more coding languages are supported in TensorFlow than in PyTorch, which has a C++ API. You can use TensorFlow in both JavaScript and Swift. If you don’t want to write much low-level code, then Keras abstracts away a lot of the details for common use cases so you can build TensorFlow models without sweating the details.

Data and Model

What models are you using? If you want to use a specific pretrained model, like BERT or DeepDream, then you should research what it’s compatible with. Some pretrained models are available in only one library or the other, and some are available on both. The Model Garden and the PyTorch and TensorFlow hubs are also good resources to check.

What data do you need? If you want to use preprocessed data, then it may already be built into one library or the other. Check the docs to see—it will make your development go faster!

Project Goal

Where will your model live? If you want to deploy a model on mobile devices, then TensorFlow is a good bet because of TensorFlow Lite and its Swift API. For serving models, TensorFlow has tight integration with Google Cloud, but PyTorch is integrated into TorchServe on AWS. If you want to enter Kaggle competitions, then Keras will let you quickly iterate over experiments.

Think about these questions and examples at the outset of your project. Nail down the two or three most important components, and either TensorFlow or PyTorch will emerge as the right choice.