TensorFlow or PyTorch? A Guide to Python Machine Learning Libraries

Let's get started

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, learn-ability, 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, FastText, NLTK)
  • visualisation 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 library 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 and want production ready solution Tensorflow 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 HardwareMachine Learning or Deep Learning?ML Knowledge required (beginner, intermediate, advanced)Learning Curve
Sci-Kit LearnLearning MLNoML onlyBeginnerVery Low
PyTorchAcademic use and productionYesBothBeginnersLow
CaffeImage processingYesBothMid-levelLow
TensorFlowProcessing large data sets quickly (production)YesBothintermediateHigh
TheanoHigh-speed computationYesBothAdvancedVery High

Deep Dive

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.


  • 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


  • 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.

import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn import svm

# load dataset containing images of digits
digits = datasets.load_digits()

# print the length of the dataset

# input vector excluding the last item to x
# Assign the target vector to y
x, y = digits.data[:-1], digits.target[:-1]

# Initiate Support Vector Machines for Classification
clf = svm.SVC(gamma=0.001, C=100)

# Fit the classifier instance to the model so that it learns from it
clf.fit(x, y)

# Predict what digit is in the last item of the input array

# Plot the image
plt.imshow(digits.images[-1], cmap=plt.cm.gray_r, interpolation="nearest")

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 Keras. The framework lets you efficiently work with mathematical expressions that include multi-dimensional arrays.


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


  • 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.

import theano
from theano import tensor

# declare two floating point scalars as a and b
a = tensor.dscalar()
b = tensor.dscalar()

# create a simple expression of multiplication
c = a * b

# convert the expression into an object that takes a and b as input
# and calculates a value for c
f = theano.function([a, b], c)

# bind 1.5 to 'a', 2.5 to 'b', and evaluate 'c'
# assert if the function, f predicts the right answer when 1.5 and 2.5 are multiplied
assert 3.75 == f(1.5, 2.5)

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!


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


  • 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.

# 1) install caffe using command `sudo apt install caffe-cpu`
# 2) train network `./examples/mnist/train_lenet.sh`
# 3) build pycaffe (instructions could be taken here - https://github.com/dungba88/caffe-python3-install/blob/master/install-caffe.md), run make pycaffe after make all

# Remarks on step 3:
# The Makefile.config file might reference to older version of python (3.5). If your system has Python 3.6, make appropriate changes (3.5 -> 3.6) in Python 3 section.
# pycaffe module builds will be located in caffe/python directory. To be able to import it you will need to set PYTHONPATH environment variable to caffe/python 
import caffe
import numpy as np
import cv2
import sys

model = "examples/mnist/lenet.prototxt"
weights = "examples/mnist/lenet_iter_10000.caffemodel"
net = caffe.Net(model, weights, caffe.TEST)

img = cv2.imread(sys.argv[1], 0)
if img.shape != [28, 28]:
    img = cv2.resize(img, (28, 28))
img = img.reshape(28, 28, -1)
img = 1.0 - img / 255.0

out = net.forward_all(data=np.asarray([img.transpose(2, 0, 1)]))

print('The recognized digit is', out["prob"][0].argmax())

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.


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


  • 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.

# Add two matrices of size 2x3 using Pytorch
import torch

a = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
a = torch.Tensor(a)
b = a[0] + a[1]
print("sum of the two matrices of the 2-D tensor:", b)

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.


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


  • 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. 

# A Basic operation using tensorflow
import tensorflow as tf

# Define 2 tensorflow constants
a = tf.constant(2, name="a")
b = tf.constant(3, name="b")
# Use the tensorflow function 'add' to add the two constants
c = tf.add(a, b, name="c")
# Tensorflow now creates a graph and runs it in a session.
sess = tf.Session()
result = sess.run(c)
# Close the session

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.


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 also find me on linkedin. I’d love to hear from you if I can help you or your team with machine learning.

This article originally appeared on Kite (Click Here)


  1. AI Patasala Python Course in Hyderabad is sure to be the ideal choice for those looking to gain insight into all the real-world challenges in this field. AI Patasala Python course is the best option for students to begin their career with the latest technology.
    Python Course in Hyderabad with Placements

  2. Know every aspect of Artificial Intelligence with Tecdecod. Here, I will share with you some unknown and secret hacks and tips of AI. There are many things you need to know about AI, and this guide will help you know everything. So, what are you waiting for? Subscribe my blog and understand the whole AI concept.

  3. This is really a good source of information, I will often follow it to know more information and expand my knowledge, I think everyone should know it, thanks. Best ftp ports service provider.

  4. Become a data science expert by joining AI Patasala’s Data Science Course in Hyderabad program, where you can learn data science concepts with practical knowledge.
    Data Science Course


Post a Comment

Popular Posts