Faster NumPy with TensorFlow. Significantly speed up your NumPy... | by Dimitris Poulopoulos | Towards Data Science

NumPy (Numerical Python) is the most popular Python library in the Data Science world. Countless other Python modules are based on NumPy, as it provides a powerful API for working with arrays and several other operations, covering branches of mathematics like linear algebra and Fourier transform. As an attribution, NumPy was created in 2005 by Travis Oliphant and it is one of the most active open-source projects.
In this story, we take NumPy a step further, increasing its execution speed using TensorFlow. This allows running NumPy code on GPU, accelerated by TensorFlow, while also allowing access to all of TensorFlow’s APIs.
But why should we use NumPy? In Python, we have lists that serve the purpose of arrays. You can always write a matrix multiplication algorithm using Python lists and a few for loops. Well, although you can do it doesn't mean you should.
NumPy aims to provide an array object that is up to 50x faster than traditional Python lists. In NumPy, this object is called ndarray and it’s the basic building block on the library. However, in this story, we take it a step further, increasing the speed of NumPy using TensorFlow.
As of TF 2.4, TensorFlow implements a subset of the NumPy API, available as tf.experimental.numpy. This allows running NumPy code on GPU, accelerated by TensorFlow, while also allowing access to all of TensorFlow's APIs. So, let’s go through the TF NumPy API step by step.
Learning Rate is my weekly newsletter for those who are curious about the world of AI and MLOps. You’ll hear from me every Friday with updates and thoughts on the latest AI news, research, repos and books. Subscribe here!
Before diving into the details, I wanted to convince you that this is something that feels like it could be worth your time. The TensorFlow NumPy implementation uses highly optimized TensorFlow kernels, that can be dispatched on GPUs and TPUs. Also, TensorFlow performs many compiler optimizations, like operation fusion, which translate to performance and memory improvements.
To compare the performance of raw NumPy against its TensorFlow counterpart, the TensorFlow team designed and ran a simple sigmoid benchmark experiment. The results in the following figure speak for themselves.
So, if these results excite you, let’s see how you can leverage this API to get similar performance.
We are ready now to dive deeper into the API, explore how TF interoperates with NumPy, but also see what’s still missing. So, what can we do with TF NumPy API?
We can accelerate NumPy code on CPU / GPU / TPU
We can auto-differentiate through NumPy code
We can run operations in a distributed manner with tf.distribute
We can interoperate seamlessly with TensorFlow APIs (tf.linalg, tf.data, tf.keras, tf.signal)
We can serialize our code using the SavedModel format and, for example, use it on edge devices with TF Lite
So, let’s see some of these points in action.
Using the TF implementation of NumPy we can place data and operations on CPUs, GPUs or even TPUs and remote devices. All we need to do is use standard TensorFlow mechanisms, such as the tf.device context manager.
The last line proves that the operations took place on the GPU, that is if we have a GPU available:
> y is placed on /job:localhost/replica:0/task:0/device:GPU:0
TensorFlow’s GradientTape can be used to enable auto-differentiation for TensorFlow NumPy code. GradientTape APIs can also return ND array outputs. To watch this in action, let’s define and instantiate a simple model:
This model is a simple two-layer neural network, with a ReLU activation between the two linearities. Its weights (w1, w2, bias) are not TF variables, but TF ND arrays. The predict method simulates a forward pass through the layers of the network.
Now we are ready to perform a forward pass through this model, calculate the loss and compute the gradients of the model, which, as we said previously, are not TF tensors, but TF ND arrays:
The last line computes the gradients of the weights of the models, which we are watching with the gradient tape. We could now perform an update, using an optimizer like the Stochastic Gradient Descent.
TensorFlow ND arrays can interoperate with NumPy functions and the other way around. However, the conversion of a TF ND array to and from a NumPy ND array may trigger actual data copies, slashing the performance.
Despite the shortcomings, this is a powerful tool to work around cases where Tensorflow doesn’t still support certain NumPy operations. Moreover, we can pass a TF ND array to functions that expect a NumPy ND array, like matplotlib’s plotting functions.
As we can see, a TF ND array can be passed to both a NumPy function (e.g., np.sum) and the plot method of matplotlib.
But what about operator precedence? Well if you add a TF ND array with a NumPy ND array, the TF one takes precedence. Also, if you pass a TF ND array to a TensorFlow function, the return type will be a TensorFlow tensor. Thus, by design, TensorFlow data types have higher precedence than those of NumPy:
Having said all that, this TensorFlow API is in an experimental phase and not every NumPy operation is supported. So, currently, TF does not support:
Mutation
Some dtypes (e.g., object, recarray)
Fortran order, views
NumPy C API, Cython and Swing integration
The tf.experimental.numpy module will be available in the stable branch starting from TensorFlow 2.4. For now, we can get it by installing the nightly distribution:
pip install --quiet --upgrade tf-nightly
In this story, we saw how we can use the new TensorFlow NumPy API, to accelerate NumPy code and how TensorFlow can now interoperate with NumPy, as well as with other libraries that are expecting np.ndarray's as input.
We also saw how to expand what you can do with NumPy, using auto-differentiation, to create and train a Machine Learning model from scratch, on GPU and using TensorFlow-related speed-up tricks, like compilation and auto-vectorization.
You can test all these now by using the nightly TensorFlow distribution and it will be available in the stable branch starting from TensorFlow 2.4.
Learning Rate is my weekly newsletter for those who are curious about the world of AI and MLOps. You’ll hear from me every Friday with updates and thoughts on the latest AI news, research, repos and books. Subscribe here!
My name is Dimitris Poulopoulos and I’m a machine learning engineer working for Arrikto. I have worked on designing and implementing AI and software solutions for major clients such as the European Commission, Eurostat, IMF, the European Central Bank, OECD, and IKEA.
If you are interested in reading more posts about Machine Learning, Deep Learning, Data Science and DataOps follow me on Medium, LinkedIn or @james2pl on Twitter.