Supercharge Python with DeepMind's JAX
Patrick Nicolas
Director Data Engineering @ aidéo technologies |software & data engineering, operations, and machine learning.
Ever hoped for numpy to offer automatic differentiation and run math computations on a GPU? You might find?DeepMind's JAX?to be the solution.?
In this piece, we'll delve into JAX's automatic differentiation capabilities and assess how its just-in-time execution compares to?numpy.
What you will learn: Applying DeepMind's JAX for speed up numerical computation and differentiation in Python.
Notes:
Introduction
As a quick recall, NumPy stands as a Python library for numerical and scientific computation. It equips data scientists and engineers with capabilities for working with multidimensional arrays, performing speedy array operations, and handling fundamental tasks in linear algebra and statistics [ref?1].
JAX [ref?2] is a numerical computing and machine learning library in Python, developed by DeepMind, that builds upon the foundation of NumPy. JAX offers:
Components
Installation
Here is an overview of basic steps for installing JAX. It is advisable to consult the installation guide as each environment has specific requirements [ref 4].
CPU (MacOS, Linux)
GPU (Linux/CUDA)
GPU (MacOS/mps)
Conda
Automatic differentiation
Overview
Automatic differentiation is a tool that facilitates the automatic calculation of derivatives for a specified mathematical function [ref :].
This technique efficiently determines precise derivatives by retaining details during the forward pass, which are then utilized during the backward pass. Essentially,?
This article focuses on the Forward Mode Automatic Differentiation, which consists of replacing each primitive operation in the original program by its differential analogue.
To illustrate the concept, let consider the function:
?Let's build its forward computation graph:
Notes:?
Single variable function
Let's implement a class, JaxDifferentiation?that wraps the computation of first, second, ... derivatives of a function with a single variable: func: R→R
.The derivative of various orders are computed in the ?constructor. The method?'__call__' return the list [f, f', f", ..].
class JaxDifferentiation(object):
"""
Create a set of derivatives of first, second, ... order_derivative order
:param func Differentiable function
:param Order of derivatives
"""
def __init__(self,
func: Callable[[float], float],
order_derivative: int):
assert order_derivative < 5,
f'Order derivatives {order_derivative} should be [0, 4]'
# Build list of derivative f, f', f", ....
self.derivatives: List[Callable[[float], float]] = [func]
temp = func
if order_derivative > 0:
for order in range(order_derivative):
# Compute the single variable next order derivative
temp = jnp.grad(temp)
self.derivatives.append(temp)
def __call__(self, x: float) -> List[float]:
""" Compute derivatives of all orders for value x"""
return [derivative(x) for derivative in self.derivatives]
Let's compute the derivative of the following function.
The first and second derivatives are provided for evaluation purpose (Oracle).
# Function definition
def func1(x: float) -> float:
return 2.0*x**4 + x**3
# First order derivative
def dfunc1(x: float) -> float:
return 8.0*x**3 + 3*x**2
# Second order derivative
def ddfunc1(x: float) -> float:
return 24.0*x**2 + 6.0*x
funcs1 = [func1, dfunc1, ddfunc1]
jax_differentiation = JaxDifferentiation(func1, len(funcs1))
compared = [f'{oracle}, {jax_value}, {oracle-jax_value}'
for oracle, jax_value in zip([func(y) for func in funcs1],
jax_differentiation(2.0))]
print(compared)
Output
Oracle, Jax, ?? Difference
40.0,??? ?40.0,?? ?0.0
76.0,??? ?76.0,?? ?0.0
108.0, 108.0,?? ?0.0
Multi-variable function
The next step is to evaluate the computation of partial derivative of a multi-variable function f(x, y,...).
Let's consider the following function for which the first order partial derivative (Jacobian vector) is provided.?
领英推荐
# Function definition
def func2(x: List[float]) -> float:
return 2.0*x[0]*x[0] - 3.0*x[0]*x[1] + x[2]
# Partial derivative over x
def dfunc2_x(x: List[float]) -> float:
return 4.0*x[0] - 3.0*x[1]
# Partial derivative over y
def dfunc2_y(x: List[float]) -> float:
return -3.0*x[0]
# Partial derivative over z
def dfunc2_z(x: List[float]) -> float:
return 1.0
Let's compare the output of the direct computation of the symbolic derivatives (Oracle)?dfunc2_x,?dfunc2_y?and?dfunc2_z?with the partial derivatives computed by JAX.
We use the?forward mode?automatic differentiation function?jacfwd?to compute the gradient [ref?6].?
# Invoke the Jacobian vector forward function
dfunc2 = jnp.jacfwd(func2)
y = [2.0, -1.0, 6.0]
derivatives = dfunc2(y)
print(f'df/dx: {derivatives[0]}, {dfunc2_x(y)}\ndf/dy: {derivatives[1]}, {dfunc2_y(y)}\ndf/dz: {derivatives[2]}, {dfunc2_z(y)}'
)
Output? ? Oracle. Jax
df/dx ?? 11.0 ?? ?11.0
df/dy ?? ? -6.0 ?? ?-6.0
df/dz. ??1.0, ?? ? ??1.0
Note: The reverse mode automatic differentiation Jax method, jacrev would have produce the same result.
Performance: JAX vs. NumPy
A significant drawback of the NumPy library is its absence of GPU support. The next objective is to measure the performance gains achieved by JAX, with and without its just-in-time compiler, on both CPU and GPU.
To facilitate this, we will establish a class named JaxNumpyData containing two functions: np_func, which utilizes NumPy, and jnp_func, its JAX counterpart. These functions will be applied to datasets of various sizes. The compare method will extract 20 subsets from the initial dataset by employing a basic fraction-based approach.
class JaxNumpyData(object):
"""
Initialize the numpy and Jax function to process data (arrays)
:param np_function Numpy numerical function
:param jnp_function Corresponding Jax numerical function
"""
def __init__(self,
np_func: Callable[[np.array], np.array],
jnp_func: Callable[[jnp.array], jnp.array]):
self.np_func = np_func
self.jnp_func = jnp_func
def compare(self, full_data_size: int, func_label: AnyStr):
"""
Compare the Numpy and JAX computation of give dataset
:param full_data_size Size of the original dataset used
to extract sub-data set
:param func_label Label used for performance results and
plotting
"""
for index in range(1, 20):
fraction = 0.05 * index
data_size = int(full_data_size*fraction)
# Execute on full_data_size*fraction elements using Numpy
x_0 = np.linspace(0.0, 100.0, data_size)
result1 = self.map_numpy(x_0, f'numpy_{func_label}')
# Execute on full_data_size*fraction elements using
# JAX and JAX-JIT
x_1 = jnp.linspace(0.0, 100.0, data_size)
result2 = self.map_jax(x_1, f'jax_{func_label}')
result3 = self.map_jif(x_1, f'jif_{func_label}')
del x_0, x_1, result1, result2, result3
"""
Process numpy array, np_x through numpy function np_func
"""
@time_it
def map_numpy(self, np_x: np.array, label: AnyStr) -> np.array:
return self.np_func(np_x)
""" Process Jax array, jnp_x through Jax function jnp_func """
@time_it
def map_jax(self, jnp_x: jnp.array, label: AnyStr) -> jnp.array:
return self.jnp_func(jnp_x)
"""
Process Jax array, jnp_x through Jax function jnp_func using JIT
"""
@time_it
def map_jif(self, jnp_x: jnp.array, label: AnyStr) -> jnp.array:
from jax import jit
return jit(self.jnp_func)(jnp_x)
The method map_numpy (resp. map_jax and map_jit) applies the NumPy method np_func (resp. JAX method jnp_func) to the NumPy array np_array (resp. JAX array jnp_array).
CPU
In this first performance test, we measure the duration to compute?
?on 1,000,000,000 values using NumPy, JAX w/o just in time compiler.
def np_func1(x: np.array) -> np.array:
return np.sinh(x) + np.cos(x)
def jnp_func1(x: jnp.array) -> jnp.array:
return jnp.sinh(x) + jnp.cos(x)
The JAX produces a 7 fold performance improvement over NumPy. The just in time processor adds another 35% improvement.
The second latency test computes the mean value:
?of a NumPy and JAX array for 1,200,000,000 values.
def np_func2(x: np.array) -> np.array:
return np.mean(x)
def jnp_func2(x: jnp.array) -> jnp.array:
return jnp.mean(x)
The just-in-time processor outperforms both NumPy and JAX native library on CPU.
GPU
For this last test, we execute the function:
over 200,000,000 values on Nvidia V100 GPU.
As anticipated, NumPy is currently running on the CPU of the EC2 instance, which means it cannot match the performance of JAX running on the Nvidia processor.
Conclusion
In summary, JAX offers data scientists and machine learning engineers a high-performance GPU computing tool that significantly outperforms the NumPy library. Our exploration has only touched the surface of JAX's capabilities, and I encourage readers to delve deeper into features like Autobatching, Vectorization, Generalized convolutions, and its integration with PyTorch and TensorFlow.
Thank you for reading this article. For more information ...
References
[1]?NumPy user guide
[4]?Installing JAX
Appendix
We include the decorator used for timing the execution of the various functions, for reference.
timing_stats = {}
def time_it(func):
""" Decorator for timing execution of methods """
def wrapper(*args, **kwargs):
start = time.time()
func(*args, **kwargs)
duration = '{:.3f}'.format(time.time() - start)
key: AnyStr = args[2]
print(f'{key}\t{duration} secs.')
cur_list = timing_stats.get(key)
if cur_list is None:
cur_list = [time.time() - start]
else:
cur_list.append(time.time() - start)
timing_stats[key] = cur_list
return 0
return wrapper
---------------------------
Patrick Nicolas has over 25 years of experience in software and data engineering, architecture design and end-to-end deployment and support with extensive knowledge in machine learning.? He has been director of data engineering at Aideo Technologies since 2017 and he is the?author of "Scala for Machine Learning" Packt Publishing ISBN 978-1-78712-238-3
#python #numpy #jax #automaticdifferentiation #vectorization #autograd #machinelearning
Assistant Professor at Faculty of Science and Techniques Tangier
1 年Interesting ??