TensorFlow or PyTorch? A Guide to Python Machine Learning Libraries (with examples!)


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.

Library Best Application Can Run on External Hardware Machine Learning or Deep Learning? ML Knowledge required (beginner, intermediate, advanced) Learning Curve
Sci-Kit Learn Learning ML No ML only Beginner Very Low
PyTorch Academic use and production Yes Both Beginners Low
Caffe Image processing Yes Both Mid-level Low
TensorFlow Processing large data sets quickly Yes Both intermediate High
Theano High-speed computation Yes Both Advanced Very 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.

Back to Top