TensorFlow vs PyTorch for Deep Learning | Towards Data Science

You may be wondering whether to learn PyTorch or TensorFlow (2.0). If so hopefully this blog post can help. The idea is not to give an absolute answer here but just to show what developing and training neural networks look like in both.
Before looking into the code, some things that are good to know: Both TensorFlow and PyTorch are machine learning frameworks specifically designed for developing deep learning algorithms with access to the computational power needed to process lots of data (e.g. parallel computing, training on GPUs, etc).
TensorFlow, which comes out of Google, was released in 2015 under the Apache 2.0 license. In Oktober 2019, TensorFlow 2.0 was released, which is said to be a huge improvement. It’s typically used in Python. PyTorch, on the other hand, comes out of Facebook and was released in 2016 under a similarly permissive open source license. As its name suggests, it’s also a Python library.
Back to the main reason for this blog post. The plan is to implement a simple neural network architecture in both TensorFlow and PyTorch to see some of the similarities and differences.
The neural network model consists of three layers: Embedding Layer → Global Average Pooling Layer → Dense Layer. It’s based on this example.
The dataset used here consists of 40,000 tweets and their sentiment (0=negative, 1=positive). First, loading the data from a CSV file and displaying some rows of the data frame to get an idea of the data.
import numpy as npimport pandas as pddata = pd.read_csv("./data/tweets.csv", encoding=’utf-8')data[:10]
To vectorize the tweets, I used Keras’ tokenizer here but there’re countless others that can do the same or even more.
Here’s what happens in the code snippet above. We assign an integer to each of the 20,000 most common words of the tweets and then turn the tweets into sequences of integers. We pad shorter ones with zeros and cut off longer ones, forcing a sequence length of 42. Finally, we should have a matrix of dimension 40,000 x 42 (tweets x sequence length).
import tensorflow as tf
Creating TensorFlow models is typically done using Keras. Keras is built on top of TensorFlow and allows for easy and fast prototyping because it has many layers built-in — it would be tedious or even prohibitive to code them from scratch each time.
There are three ways to build a neural network model in Keras.
You can create your own fully-customizable models by subclassing the tf.keras.Model class and implementing the forward pass in the call method. Put differently, layers are defined in the __init__() method and the logic of the forward pass in the call method.
The benefit of using such an object-oriented approach is that you can reuse layers multiple times within the call method or define a more complex forward pass. But nothing like that is happening in this example here, it’s just a linear stack of layers.
In the functional API, given some input tensor(s) and output tensor(s), you can also instantiate a Model. It’s a user-friendly way to build a neural network and Keras even recommends it over model subclassing. With this approach, you essentially define a layer and immediately pass it the input of the previous layer. So it requires slightly less coding with the same result.
The sequential API is the most compact way to define a model and sufficient for certain (simple) neural networks, typically consisting of just a few common layers — kind of a shortcut to a trainable model. It’s definitely convenient and works well but is too inflexible if you wish to implement more sophisticated ideas.
Regardless of how you build a Keras model, there are two functions I’d like to show you. First, calling model.summary() prints a compact summary of the model and the number of parameters, which is super useful.
Second, by calling tf.keras.utils.plot_model() you get a graphical summary of the model. It didn’t work in my Jupyter Notebook, however. So I can’t show it to you here, unfortunately.
Before you can train a Keras model, it must be compiled by running the model.compile() function, which is also where you specify the loss function and optimizer.
model.compile(loss='binary_crossentropy', optimizer='Adam', metrics=['accuracy'])
Keras models have a convenient fit function for training a model (just like Scikit-Learn), which also takes care of batch processing and even evaluates the model on the run (if you tell it to do so).
model.fit(x=X, y, batch_size=32, epochs=5, verbose=2, validation_split=0.2)
Note: It’s okay to pass Numpy arrays as inputs to the fit function even though TensorFlow (PyTorch too for that matter) operates on tensors only, which is a similar data structure but optimized for matrix computations. Keras takes care of transforming the arrays under the hood.
import torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.optim as optim
There are two ways to build a neural network model in PyTorch.
Similar to TensorFlow, in PyTorch you subclass the nn.Model module and define your layers in the __init__() method. The only difference is that you create the forward pass in a method named forward instead of call.
Note: Difference to the Keras model: There’s only an average-pooling layer in PyTorch so it needs to have the right kernel size in order the make it global average-pooling.
PyTorch also offers a Sequential module that looks almost equivalent to TensorFlow’s.
Note: I found that many layers do not work with PyTorch’s nn.Sequential such as many recurrent layers (RNNs, LSTMS, etc.). In fact, PyTorch didn’t really want to implement the sequential module at all because it wants developers to use subclassing.
There’s no pre-made fit function for PyTorch models, so the training loop needs to be implemented from scratch. Here’s what a typical training loop in PyTorch looks like.
Note: In oder to process the data in batches, a dataloader must be created (left out here, see GitHub code). The dataloader returns one batch at a time in a dictionary format.
Short description of the training loop: For each batch, we calculate the loss and then call loss.backward() to backpropagate the gradient through the layers. In addition, we call optimizer.step() to tell the optimizer to update the parameters. For a more detailed description of how to train a PyTorch model see here.
Here’re the best resources I know about getting started with TensorFlow and/or PyTorch.
PyTorch’s official tutorial on their website is awesome and in my view better than TensorFlow’s. There’s also a really good free PyTorch course on Udacity.
Tutorials on the PyTorch Website
Intro to Deep Learning with PyTorch on Udacity
Deep Learning with PyTorch (Book)
TensorFlow has some really good courses on Coursera, and there’s a great book written by François Chollet, the creator of Keras. TensorFlow’s and especially Keras’ official websites are important sources too.
TensorFlow Website
Keras Dokumentation
Intro to TensorFlow for Deep Learning on Udacity
TensorFlow in Practice from deeplearning.ai on Coursera
Deep Learning with Python by François Chollet (Book)
I don’t really feel able to answer that because I have only scratched the surface so far. But here are some of the things I’ve noticed.
TensorFlow is a lot like Scikit-Learn thanks to its fit function, which makes training a model super easy and quick. But even though you need to build every training loop in PyTorch yourself, I kind of like it because it makes you think more carefully about what you’re doing. So it’s less of a black box and that’s always good. I fist learned PyTorch and I think that was a very good idea.
Note: Of course, TensorFlow allows you to build customized training loops as well, but it’s not as neat. Plus, there are third-party libraries for PyTorch that automate the training loop so that the difference here may not be so big after all.
All in all, it’s certainly easier to go from a blank script to a trained neural network in TensorFlow than in PyTorch — mostly due to TensorFlow’s fit method. But another difference plays a role here too. I’ve noticed that Keras layers often do not require you to specify the input dimension whereas in PyTorch you need to be more explicit. It’s easy to confuse the dimensions, especially if you want to try out different ones. Apart from that, Keras’ model.summary() and tf.keras.utils.plot_model()are super useful functions, as mentioned before. There are no standard PyTorch functions (I’m aware of) that can do the same.
These are definitely some plus points for TensorFlow. But in general, the two libraries feel VERY similar. It seems that they have converged a lot already by learning from each other and adopted each other’s best features. That’s not at all a bad thing but simply a sign that they have matured a lot. Therefore it’s also really easy to switch — so don’t worry about choosing the “wrong” library. For sure, there are a lot more aspects that I did not consider here, especially any advanced features such as parallel computing, training on GPUs, etc.
Anyway, it will be interesting to see how TensorFlow and PyTorch will do in 2020. While PyTorch has been more popular among researchers lately, TensorFlow is the frontrunner in the industry. Thanks to TensorFlow and PyTorch, deep learning is more accessible than ever and more people will use it. It’s never been easier.
The code can also be found as a Jupyter Notebook here.