TensorFlow or PyTorch? A guide to Python machine learning libraries

by Kite 26 September 2019

Python is the fastest-growing programming language out there. That isn’t surprising given that it’s simple, easy to use, free, and applicable for many computing tasks. Data scientists in particular have embraced Python’s efficient syntax, learnability, and easy integrations with other languages such as C and C++.

All these positive qualities, along with the recent spike of interest in machine learning and artificial intelligence, can help explain the plethora of powerful open-source libraries and frameworks for machine learning and data science applications. There are libraries that can be put to use in a multitude of applications, including:

  • natural language processing / NLP (Tensorflow)
  • visualization and analysis of complex data (Theano)
  • image recognition (Caffe)
  • prediction and recommendation

Open-source frameworks have popped up to address all of the above applications, and now it can be confusing to decide on which library to use for which project. Tensorflow or Sci-kit? Should I use Keras on top of Microsoft’s CNTK? What’s the best application to use MXNet?

Once you’ve determined the goals and overall priorities for your project, this article can help you select the language that is the best fit for your project. Some of the questions that you’ll need to consider include:

  • Your confidence level with machine learning fundamentals
  • If you will be using the framework for classic machine learning algorithms or for Deep Learning
  • What application you will be using the framework for: be it heavy numerical computations, complex data analysis, image analysis, or education and research
  • Whether or not you’ll be using any additional hardware (like GPUs and TPUs), software, or cloud services for scaling on to bigger data sets

Each open-source framework available today has its own strengths and weaknesses when measured across these factors. And choosing the best framework for your needs will really depend on just what you want to accomplish.

For example, if you are new to machine learning or want to use classic machine learning algorithms, Sci-kit could be the best choice. On the other hand, if you need to do heavy numerical computations, Theano would work much better. In any case, no matter your specific situation – this guide will aim to help you figure out which framework is the perfect fit.

LibraryBest applicationCan run on external hardwareML or DL?ML knowledge required Learning curve
Sci-Kit LearnLearning MLNoMLBeginnerVery Low
PyTorchAcademic use and productionYesBothBeginnersLow
CaffeImage processingYesBothMid-levelLow
TensorFlowProcessing large data sets quicklyYesBothIntermediateHigh
TheanoHigh-speed computationYesBothAdvancedVery High

Among all the myriad of options available for open-source Python frameworks, here is the compilation of our top 5 choices in descending order. You can follow along with examples for each library, stored in Kite’s github repository.

5. Sci-Kit Learn

Ideal for: ML beginners

Sci-kit Learn is a library that features a host of the classical machine learning algorithms like Support Vector Machines (SVMs), KNN Maps, K-Nearest Neighbors (KNN) classifiers, Random Forests, and regression algorithms. It includes options for both supervised and unsupervised learning. Thus, it’s ultimately an effective tool for statistical modeling.

It has been built on many other Python libraries like SciPy, Numpy, and Matplotlib, and some of its core algorithms are also written using Cython. I created an example of a Sci-Kit operation here.

Strengths:

  • Great for beginners and for those looking to explore machine learning algorithms
  • Good for data-mining and simple projects like predictions on small or labeled data sets

Weaknesses:

  • Does not support ANNs
  • Does not support GPU computing

What sets this framework apart from others is an easy-to-use interface for developers and a high level of abstraction that allows especially beginners in machine learning to get easily acquainted with the platform, without having to deal with the nitty-gritty of actual algorithms.

It’s easy to run and debug, and there are some nice and easy tutorials available to help understand the algorithms when you do have to work with them. However, Sci-kit Learn does have a couple of limitations.

First, it does not support Artificial Neural Networks.

Second, it’s only suitable for small projects with small datasets, and for tasks that are not particularly computationally intensive. This is mainly due to the fact that the framework does not support GPU computing.

For more seasoned or hard-core developers, it can feel limiting to some extent, as the abstraction doesn’t allow for fine tuning the underlying algorithms.

4. Theano

Ideal for: Hardcore developers requiring high-speed computation over a single GPU

If you’re looking for a framework that can crunch numbers like a pro, then Theano will be your best bet.

Theano is a workhorse well-equipped for numerical computing and sits under a large number of other deep learning frameworks like Tensorflow and Keras. The framework lets you efficiently work with mathematical expressions that include multi-dimensional arrays.

Strengths:

  • Efficiency in crunching large and multi-dimensional data sets
  • Provides the developer ample flexibility to fine-tune underlying algorithms and create novel models.

Weakness:

  • A very steep learning curve
  • Does not support scaling over multiple GPUs

While Sci-kit Learn is for beginners, Theano is only for advanced deep learning experts.

The API is low-level, so you really need feel comfortable in your coding abilities if you’re looking to explore this framework. The syntax for Theano is quite tightly integrated with NumPy and its code can run efficiently – both on a CPU and a GPU. You can see an example of a script using Theano here.

In contrast to Sci-kit learn, Theano empowers any developer with a complete flexibility to fine-tune and control their models. It even allows the developer to implement and test completely unconventional models.

Although Theano works better than Tensorflow over a single GPU, it still doesn’t match up to Tensorflow when working with multiple GPUs.

However, since Theano has been around longer, it does have a lot more documentation. The biggest drawback is that MILA, the makers of Theano, have decided to stop maintaining the framework following its 1.0 release. Nonetheless, it continues to be a great choice for avid deep learning enthusiasts.

3. Caffe

Ideal for: Mid-level programmers and image processing

Caffe (Convolutional Architecture for Fast Feature Embedding) was mainly built to support Convolutional Neural Networks (CNNs) and is the framework of choice for those working on computer vision, image processing, and feedforward networks.

The framework is a Python-based API, which was mainly written in C++. Models in Caffe are represented by Protobuf configuration files and the framework, is, in fact, the fastest CNN implementation among all Deep Learning frameworks.

It works well on image segmentation and classification tasks. With a single GPU, Caffe can process more than 60 million images in a day! Here is a simple example of using a pre-trained Caffe model to correctly identify an image as the digit, “5”.

Strengths:

  • Has great ready-to-use models for image recognition
  • It is the fastest CNN implementation framework
  • Models and optimizations are configured rather than coded

Weaknesses:

  • Not suitable for RNNs
  • Poor documentation
  • Creating new layers requires defining full forward, backward and gradient updates

The best thing about Caffe is that models and optimizations are not ‘coded’, but rather ‘configured’ – this reduces a lot of headaches for developers.

So, if you have a large set of images that you need to classify or run some regression algorithms on, you can quickly apply a DL Network without having to write even a line of code.

You can also train over multiple GPUs, but this has some limitations: for example, you can’t do model/data parallelism.

Another upside of Caffe is that it has a pretty good developer community – there’s a whole ‘Model Zoo’ available where you can find a number of CNN implementations and models, like AlexNet, GoogleNet, NIN, etc. This is something that other frameworks lack.

Although the framework is great for CNNs and image processing, it’s not suitable for Recurrent Neural Networks (RNNs) and applications involving text, sound, and time series data. Moreover, even though there are a lot of layers ready to be implemented, the creation of new layers can be tedious, as one would need to define full forward, backward and gradient updates for each new layer.

Finally, the framework offers a medium-level abstraction – it’s high-level enough to allow you to do quick experiments and flexible enough to allow you to fine-tune some of the aspects. This detail may be a positive for mid-level developers, but it feels somewhat limiting for hardcore developers.

2. Pytorch

Ideal for: Both academic use and production

Pytorch was developed using Python, C++ and CUDA backend. Created by the Facebook Artificial Intelligence Research team (FAIR), Pytorch is fairly new but is already competing neck-to-neck with Tensorflow, and many predict it will soon become a go-to alternative to many other frameworks.

Strengths:

  • Coding is easy, so it has a flatter learning curve
  • Supports dynamic graphs so you can adjust on-the-go.
  • Supports GPU acceleration

Weaknesses:

  • Quite new, so it has a smaller community and fewer resources available online

Pytorch is being lauded particularly by beginners, mostly due to its easy-to-write code – but the framework is basically a blend of both high and low-level APIs. In actuality, it’s suitable for both academic uses as well as hard-core deep learning.

It features a number of pre-trained models. When coding in Pytorch, you don’t need to categorize numbers into ‘int’, ‘short’, or ‘double’ data types, like other coding languages. This makes the performance of operations and functions on this framework more intuitive compared to other options. You can see example code for PyTorch here.

The highlight of this framework, though, is that it offers developers the ability to use dynamic graphs. Tensorflow, Theano, and their derivatives allow you to create only static graphs, so you have to define the whole graph for the model before you can run it. However, in Pytorch, you can define or adjust your graph during runtime, so it’s more flexible and allows you to use variable length inputs, especially in your RNNs.

The framework also provides strong support for GPU acceleration, so you get both efficiency and speed.

The main drawback to Pytorch, though, is that the framework is still growing and you may encounter some bugs. Moreover, owing to its younger age, the resources to supplement its official documentation are still quite scant. But looking at overall trends, this will not be a problem for too long, as more and more developers are converting to Pytorch and the community is growing slowly but steadily.

1. Tensorflow

Ideal for: Intermediate-level developers and for developing production models that need to quickly process vast data sets

Tensorflow is currently hailed as the best ML framework out there. Within a very short time, it has become a favorite for many developers and is witnessing an ever-growing community and extraordinary development momentum.

The framework was developed by the Google Brain team and supports all platforms, from Linux to Android. It is a high-level framework that allows you to run low-level code with supporting libraries. Ultimately, it allows you to monitor the progress of the training process, while tracking a lot of metrics and not having to bother about most of the other details.

Strengths:

  • Flexibility
  • Contains several ready-to-use ML models and ready-to-run application packages
  • Scalability with hardware and software
  • Large online community

Weaknesses:

  • Supports only NVIDIA GPUs
  • A slightly steep learning curve

Tensorflow’s architecture and UX are different from other frameworks in that the nodes in a Tensorflow graph represent mathematical operations, while the edges of the graph represent multidimensional arrays (tensors). These tensors flow between the nodes, giving you a lot of flexibility when it comes to creating new nodes, unlike the Caffe architecture, for example.

The system also has a host of models to choose from: the framework is pre-loaded with packages that let you perform voice recognition and machine translation, and models that let you run regressions, classifications, neural networks and an assortment of other algorithms.

Tensorflow can be used for quite a few applications within machine learning. Check out a basic “Hello, World” program here and a more traditional matrix example here.

But the feature that really takes the cake is Tensorflow’s computing capabilities. To date, Tensorflow is the strongest contender in the distributed processing arena. It provides remarkable scalability and lets you deploy your computations to multiple CPUs, GPUs, other servers, mobile devices, and the Google Cloud Machine Learning Engine. You can do this without having to rewrite any code – that’s truly powerful.

The main downside, though, is that at the moment, it only supports NVIDIA GPUs. Also, when it comes to RNN support, it is ultimately weaker than some other frameworks and the learning curve can be a little steeper than Sci-kit and Pytorch.

Overall, with a strong Google backing and a huge online community, Tensorflow is here for the long haul.

Conclusions

To sum up, while Tensorflow has gained enormous popularity owing to its flexibility and distributed processing capabilities, Pytorch is also slowly gaining momentum owing to its flatter learning curve and ability to process dynamic graphs. Both of these frameworks are multi-purpose and can be applied to many types of projects. Other frameworks like Caffe, Theano, and Sci-Kit Learn are more specialized and aimed toward specific tasks.

Nonetheless, machine learning and artificial intelligence are the future, and these open source frameworks have brought ML within the grasp of any developer with a really keen interest. These libraries provide the tools for any Pythonista to practice machine learning principles. Pick a framework that meets your level of expertise and application domain and try out your first project!


You can check out the code from this and other posts on Kite’s GitHub repository.

This article originally appeared on Kite. Want to code faster? Kite is a plugin for PyCharm, Atom, Vim, VSCode, Sublime Text, and IntelliJ that uses machine learning to provide you with code completions in real time sorted by relevance.