Which Framework to Choose: TensorFlow vs. PyTorch

Which Framework to Choose: TensorFlow vs. PyTorch

What are we talking about?

The rapid development of deep learning (DL), which is part of the field of artificial intelligence (AI) and is based on artificial neural networks, has led to strong competition between DL frameworks and passionate debates. In recent years, two clear leaders have emerged: TensorFlow by Google and PyTorch by Facebook. Initially, TensorFlow (TF) took the lead in the industry, strengthening its position after Google hired François Chollet, a great AI researcher and author of another popular framework called Keras. Then came the new kid on the block, PyTorch, and the situation began to change, TensorFlow now had a formidable competitor with a rapidly growing adoption rate. Naturally, many companies and AI teams were faced with a tough choice of which framework to use. This of course also resulted in a never-ending list of “my framework is better than yours” debates on Reddit (& Twitter).  In this post, I am going to share my thoughts on this topic based on my and my team’s practical experience with both frameworks.

Current state

First of all, today both frameworks have large communities (TFPyTorch) and, feeding of each other's success and competition, are gradually adding more and more features. As a result, nowadays both frameworks have very similar capabilities. Soumith Chintala, one of the co-creators of PyTorch also thinks so and earlier this year he made his feelings very clear with this Tweet:


Soumith Chintala Tweet

From my experience

Based on my experience, here are the key points you need to be aware of when it comes to TensorFlow vs PyTorch:

  • TF can win in speed. This primarily applies to the models created by Google (for example, BERT) and to the family of TF libraries. As an illustration, I recently studied two different implementations of the finetuning of the BERT base:
    • 1) PyTorch + transformers, and
    • 2) using tensorflow-hub

The latter was faster by 15-20% and worked with the batch size 8 (for PyTorch + transformers, the max batch size for the same GPU was 6).

  • With TF, it is easier to work with Google's TPU, but it is important to note PyTorch Lightning is showing a lot of promise in this area.

  • In my opinion, the TF code is much easier to read due to the domination of Keras in TF 2.x. But this is subjective, simply a personal preference.

  • In production, TF is still ahead of PyTorch due to TF Lite, TF Serving, etc. Of course, PyTorch is catching up (and PyTorch Lightning is good evidence of that), but would Google stand still?

  • Finally, as a professor, I would like to note that Google is very active in creating training materials for TF, which simplifies the adoption of the framework. Particularly successful for Google is the partnership with Andrew Ng’s DeepLearning.AI, which has resulted in such online specializations as TensorFlow in Practice (now – DeepLearning.AI TensorFlow Developer Professional Certificate), TensorFlow: Data and Deployment, and the newest TensorFlow: Advanced Techniques. However, it is worth noting, that the recent Generative Adversarial Networks (GANs) Specialization by DeepLearning.AI is based on PyTorch.

It must be noted, that despite my personal preferences and views, relying solely on TF is not a robust strategy. With such giants as Facebook, Microsoft, Amazon, OpenAI, Tesla, Uber, etc. having united around PyTorch, the active development & advancement of the framework is very much guaranteed.

Therefore, I do not think we should focus too much on just one framework, especially when it comes to research tasks. I will personally continue to improve my knowledge of both frameworks and recommend this approach to anyone who is serious about a career in AI. In R&D activities, you must be flexible and adaptable to be able to compare and choose the best solution available for the current problem you need to solve. Also, the task has been greatly simplified now as the field has been narrowed down from over a dozen different frameworks to just two… so let’s stop debating and let’s get to learning :)