If you have functions that do a lot of mathematical operations, use NumPy or rely heavily on loops, then there is a way to speed them up significantly with one line of code. Ok, two lines if you count the import.
Numba and the @jit decorator
Meet Numba and its @jit decorator. It changes how your code is compiled, often improving its performance. You don’t have to install any special tools (just the
numba pip package), you don’t have to tweak any parameters. All you have to do is:
- Add the
@jitdecorator to a function
- Check if it’s faster
Let’s see an example of code before and after applying
# numba_testing.py import math def compute(): # Bunch of dummy math operations result = 0 for number in range(1_000_000): double = number * 2 result += math.sqrt(double) + double return result
The only purpose of this code is to do some calculations and to “be slow.” Let’s see how slow (benchmarks are done with Python 3.8 - I describe the whole setup in the Introduction article):
$ python -m timeit -s "from numba_testing import compute" "compute()" 1 loop, best of 5: 217 msec per loop
Now, we add
@jit to our code. The body of the function stays the same, and the only difference is the decorator. Don’t forget to install Numba package with pip (
pip install numba).
# numba_testing.py import math from numba import jit @jit def compute_jit(): # Bunch of dummy math operations result = 0 for number in range(1_000_000): double = number * 2 result += math.sqrt(double) + double return result
Let’s measure the execution time once more:
$ python -m timeit -s "from numba_testing import compute_jit" "compute_jit()" 200 loops, best of 5: 1.76 msec per loop
Using @jit decorator gave us a 120x speedup (217 / 1.76 = 123.295)! That’s a huge improvement for such a simple change!
How did I find Numba?
I first learned about Numba when I was doing code challenges from the Advent of Code a few years ago. I wrote a pretty terrible algorithm, left it running, and went for lunch. When I came back after one hour, my program wasn't even 10% done. I stopped it, added the
@jit decorator to the main function, rerun it, and I had the results in under one minute! Fantastic improvement with almost no work!
This story doesn't mean that it's ok to write sloppy code, and then use hacks to speed it up. But sometimes you just need to make some one-off calculations. You don't want to spend too much time writing the perfect algorithm. Or maybe you can't think of a better algorithm, and the one you have is too slow. Using tools like Numba can be one of the fastest and easiest to apply improvements!
Other features of Numba
@jit is the most common decorator from the Numba library, but there are others that you can use:
- @njit - alias for @jit(nopython=True). In
nopythonmode, Numba tries to run your code without using the Python interpreter at all. It can lead to even bigger speed improvements, but it’s also possible that the compilation will fail in this mode.
- @vectorize and @guvectorize - produces
ufuncused in NumPy.
- @jitclass - can be used to decorate the whole class.
- @cfunc - declares a function to be used as a native callback (from C or C++ code).
There are also advanced features that let you, for example, run your code on GPU with @cuda.jit. This doesn’t work out of the box, but it might be worth the effort for some very computational-heavy operations.
Numba has plenty of configuration options that will further improve your code’s execution time if you know what you are doing. You can:
- Disable GIL (Global Interpreter Lock) with
- Cache results with
- Automatically parallelize functions with
Numba is a great library that can significantly speed up your programs with minimal effort. Given that it takes less than a minute to install and decorate some slow functions, it’s one of the first solutions that you can check when you want to quickly improve your code (without rewriting it).
It works best if your code:
- Uses NumPy a lot
- Performs plenty of mathematical operations
- Performs operations is a loop