Memoryview Benchmarks
There was recently a thread on cython-users which caught my eye. It has to do with memoryviews, a new way of working with memory buffers in cython.
I've been thinking recently about how to do fast and flexible memory buffer access in cython. I contributed the BallTree implementation for nearest neighbors searching in scikit-learn, and have been actively thinking about how to make it faster and more flexible, including adding the ability to specify distance metrics other than euclidean and minkowski.
In order to accomplish this, I'd like to have a set of distance metric functions which take two vectors and compute a distance. There would be many functions with similar call signatures which could then be plugged into a code that would iterate over a set of vectors and compute the appropriate distances.
Pure python version
In pure python, the implementation described above might look something like this:
# memview_bench_v1.py
import numpy as np
def euclidean_distance(x1, x2):
x1 = np.asarray(x1)
x2 = np.asarray(x2)
return np.sqrt(np.sum((x1 - x2) ** 2))
This looks promising. Let's create a function based on this which will compute the pairwise distance between all points in a matrix (this is similar to pairwise_distances in scikit-learn or pdist in scipy). The simple form of the function might look like this:
# memview_bench_v1 (continued)
def pairwise(X, metric=euclidean_distance):
X = np.asarray(X)
n_samples, n_dim = X.shape
D = np.empty((n_samples, n_samples))
for i in range(n_samples):
for j in range(n_samples):
D[i, j] = metric(X[i], X[j])
return D
We could exploit symmetry to reduce the number of computations required, but
we'll skip that step for now: this simple version of the function will give
us a good benchmark for comparison with alternatives below. Using the
timeit magic in ipython, we can learn how fast this implementation is:
In [1]: import numpy as np
In [2]:from memview_bench_v1 import pairwise
In [3]: X = np.random.random((500, 3))
In [4]: timeit pairwise(X)
1 loops, best of 3: 6.51 s per loop
It takes nearly seven seconds to compute 250,000 distances. This is much too slow.
Cython Speedup
Perhaps we can speed this up using cython declarations. Before typed memoryviews were added in cython 0.16, the way to quickly index numpy arrays in cython was through the numpy specific syntax, adding type information to each array that specifies its data type, its dimension, and its order:
# memview_bench.pyx
import numpy as np
cimport numpy as np
from libc.math cimport sqrt
cimport cython
# define a function pointer to a metric
ctypedef double (*metric_ptr)(np.ndarray, np.ndarray)
@cython.boundscheck(False)
@cython.wraparound(False)
cdef double euclidean_distance(np.ndarray[double, ndim=1, mode='c'] x1,
np.ndarray[double, ndim=1, mode='c'] x2):
cdef double tmp, d
cdef np.intp_t i, N
d = 0
N = x1.shape[0]
# assume x2 has the same shape as x1. This could be dangerous!
for i in range(N):
tmp = x1[i] - x2[i]
d += tmp * tmp
return sqrt(d)
@cython.boundscheck(False)
@cython.wraparound(False)
def pairwise(np.ndarray[double, ndim=2, mode='c'] X not None,
metric = 'euclidean'):
cdef metric_ptr dist_func
if metric == 'euclidean':
dist_func = &euclidean_distance
else:
raise ValueError("unrecognized metric")
cdef np.intp_t i, j, n_samples
n_samples = X.shape[0]
cdef np.ndarray[double, ndim=2, mode='c'] D = np.empty((n_samples,
n_samples))
for i in range(n_samples):
for j in range(n_samples):
D[i, j] = dist_func(X[i], X[j])
return D
Notice that we're essentially running the same code, except we have added
type identifiers to speed up function calls and loops. The mode='c'
argument in the np.ndarray type says that the array is contiguous in
memory, and C-ordered.
For reference, this can be compiled in-place by running
python setup.py build_ext --inplace with the following
setup.py file:
# setup.py
import os
import numpy
from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
module = 'memview_bench_v2'
setup(cmdclass = {'build_ext': build_ext},
name=module,
version='1.0',
ext_modules=[Extension(module,
[module + ".pyx"])],
include_dirs=[numpy.get_include(),
os.path.join(numpy.get_include(), 'numpy')]
)
We'll time the resulting function on the same sized array as we did previously:
In [1]: import numpy as np
In [2]: from memview_bench_v2 import pairwise
In [3]: X = np.random.random((500, 3))
In [4]: timeit pairwise(X)
1 loops, best of 3: 668 ms per loop
That's a factor of 10 speedup over the pure python version! It turns out,
though, that we can do better. In particular, the slicing operation when
we call X[i] and X[j] must generate a new numpy array each time, which
leads to a lot of python overhead in reference counting, etc. This is the
reason that the cython team introduced typed memoryviews in cython v0.16.
Typed Memoryviews
The equivalent of the above code using typed memoryviews looks like this:
# memview_bench_v3.pyx
import numpy as np
cimport numpy as np
from libc.math cimport sqrt
cimport cython
# define a function pointer to a metric
ctypedef double (*metric_ptr)(double[::1], double[::1])
@cython.boundscheck(False)
@cython.wraparound(False)
cdef double euclidean_distance(double[::1] x1,
double[::1] x2):
cdef double tmp, d
cdef np.intp_t i, N
d = 0
N = x1.shape[0]
# assume x2 has the same shape as x1. This could be dangerous!
for i in range(N):
tmp = x1[i] - x2[i]
d += tmp * tmp
return sqrt(d)
@cython.boundscheck(False)
@cython.wraparound(False)
def pairwise(double[:, ::1] X not None,
metric = 'euclidean'):
cdef metric_ptr dist_func
if metric == 'euclidean':
dist_func = &euclidean_distance
else:
raise ValueError("unrecognized metric")
cdef np.intp_t i, j, n_samples
n_samples = X.shape[0]
cdef double[:, ::1] D = np.empty((n_samples, n_samples))
for i in range(n_samples):
for j in range(n_samples):
D[i, j] = dist_func(X[i], X[j])
return D
The only change is that instead of using the np.ndarray[...] type specifier,
we use the typed memoryview double[:, ::1] specifier. The ::1 in the
second position means that we are passing a two-dimensional array, which
is contiguous and C-ordered. We time the results and see the following:
In [1]: import numpy as np
In [2]: from memview_bench_v3 import pairwise
In [3]: X = np.random.random((500, 3))
In [4]: timeit pairwise(X)
10 loops, best of 3: 22 ms per loop
This gives another factor of 30 improvement over the previous version, simply by switching to typed memoryviews rather than the numpy interface. Still, our function is creating memoryview objects each time we slice the array. We can determine how much overhead this is generating by using raw C pointers instead. It is not as clean, but it should be very fast:
Raw Pointers
The fundamental benchmark for this sort of operation should be working directly with the pointers themselves. While this is not a very "pythonic" way of doing things, it does lead to very fast code, as we will see:
# memview_bench_v4.pyx
import numpy as np
cimport numpy as np
from libc.math cimport sqrt
cimport cython
# define a function pointer to a metric
ctypedef double (*metric_ptr)(double*, double*, int)
@cython.boundscheck(False)
@cython.wraparound(False)
cdef double euclidean_distance(double* x1,
double* x2,
int N):
cdef double tmp, d
cdef np.intp_t i
d = 0
for i in range(N):
tmp = x1[i] - x2[i]
d += tmp * tmp
return sqrt(d)
@cython.boundscheck(False)
@cython.wraparound(False)
def pairwise(double[:, ::1] X not None,
metric = 'euclidean'):
cdef metric_ptr dist_func
if metric == 'euclidean':
dist_func = &euclidean_distance
else:
raise ValueError("unrecognized metric")
cdef np.intp_t i, j, n_samples, n_dim
n_samples = X.shape[0]
n_dim = X.shape[1]
cdef double[:, ::1] D = np.empty((n_samples, n_samples))
cdef double* Dptr = &D[0, 0]
cdef double* Xptr = &X[0, 0]
for i in range(n_samples):
for j in range(n_samples):
Dptr[i * n_samples + j] = dist_func(Xptr + i * n_dim,
Xptr + j * n_dim,
n_dim)
return D
Instead of passing around slices of arrays, we've accessed the raw memory
buffer using C pointer syntax. This is not as easy to read, and can lead
to glibc errors or segmentation faults if we're not careful. Testing
this implementation, we find that it is extremely fast:
In [1]: import numpy as np
In [2]: from memview_bench_v4 import pairwise
In [3]: X = np.random.random((500, 3))
In [4]: timeit pairwise(X)
100 loops, best of 3: 2.47 ms per loop
This is another factor of 10 faster than the memoryview benchmark above! Essentially, what this is telling us is that creating a memoryview slice takes about 0.02 / 500,000 = 40 nanoseconds on our machine. This is extremely fast, but because we're performing this operation half a million times, the cost of the allocations is significant compared to the rest of our computation. If our vectors were, say, length 1000, this cost may not be a significant percentage of the total cost.
So what are we left with? Do we need to use raw pointers in all circumstances when working with collections of small vectors? Perhaps not.
A Faster Implementation with Memoryviews
The creation of memoryview slices, though extremely fast, is causing a problem simply because we're creating so many slices. Here is an alternative which uses no raw pointers, but matches the speed of raw pointers:
# memview_bench_v5.pyx
import numpy as np
cimport numpy as np
from libc.math cimport sqrt
cimport cython
# define a function pointer to a metric
ctypedef double (*metric_ptr)(double[:, ::1], np.intp_t, np.intp_t)
@cython.boundscheck(False)
@cython.wraparound(False)
cdef double euclidean_distance(double[:, ::1] X,
np.intp_t i1, np.intp_t i2):
cdef double tmp, d
cdef np.intp_t j
d = 0
for j in range(X.shape[1]):
tmp = X[i1, j] - X[i2, j]
d += tmp * tmp
return sqrt(d)
@cython.boundscheck(False)
@cython.wraparound(False)
def pairwise(double[:, ::1] X not None,
metric = 'euclidean'):
cdef metric_ptr dist_func
if metric == 'euclidean':
dist_func = &euclidean_distance
else:
raise ValueError("unrecognized metric")
cdef np.intp_t i, j, n_samples, n_dim
n_samples = X.shape[0]
n_dim = X.shape[1]
cdef double[:, ::1] D = np.empty((n_samples, n_samples))
for i in range(n_samples):
for j in range(n_samples):
D[i, j] = dist_func(X, i, j)
return D
Timing this implementation we find the following:
In [1]: import numpy as np
In [2]: from memview_bench_v5 import pairwise
In [3]: X = np.random.random((500, 3))
In [4]: timeit pairwise(X)
100 loops, best of 3: 2.45 ms per loop
Just as fast as using raw pointers, but much cleaner and easier to read.
Summary
Here are the timing results we've seen above:
- Python + numpy: 6510 ms
- Cython + numpy: 668 ms
- Cython + memviews (slicing): 22 ms
- Cython + raw pointers: 2.47 ms
- Cython + memviews (no slicing): 2.45 ms
So what have we learned here? First of all, typed memoryviews are fast. Blazing fast. If used correctly, they can be comparable to raw pointers, but are much cleaner easier to debug. For example, in the last version, if we ran into a memory error we could simply turn on bounds-checking and quickly find the source of the problem. Slicing with memoryviews is also fast, but should be used carefully if your operation time on each slice is compararable to the cost of building the slice.
The moral of the story? Use typed memoryviews. It will lead to fast cython code which is cleaner, more readable, and more easily debuggable than any other alternative.
Update: All of the above scripts are now available as an ipython notebook: memview_bench.ipynb. For information on how to view this file, see the IPython page Alternatively, you can view this notebook (but not modify it) using the nbviewer here.
Thanks to Dave for the tip.