EDIT: There was a bug in the final version of the code presented here. It is fixed now, for its backstory, check out my blog post on it.
When we last saw our hero, he was fighting with the dreaded implementation of least-angle regression, knowing full well that it was his destiny to be faster.
We had come up with a more robust implementation, catching malformed cases that would have broken the naive implementation, and also it was orders of magnitude faster than said implementation. However, our benchmark [1] showed that it was a couple of times slower than least-angle regression.
By poking around the scikits.learn codebase, I noticed that there is a
triangular system solver in scikits.learn.utils.arrayfuncs. Unlike the
scipy.linalg one, this one only works with lower triangular arrays,
and it forcefully overwrites b. Even though if weren’t faster, it
should still be used: scikits.learn aims to be as backwards-compatible
with SciPy as possible, and linalg.solve_triangular was added in
0.9.0. Anyway, let’s just see whether it’s faster:
[sourcecode language=”python”]
In 1: import numpy as np
In 2: from scipy import linalg
In [3]: from scikits.learn.datasets import make_spd_matrix
In [4]: from scikits.learn.utils.arrayfuncs import solve_triangular
In [5]: G = make_spd_matrix(1000)
In [6]: L = linalg.cholesky(G, lower=True)
In [7]: x = np.random.randn(1000)
In [8]: y = x.copy()
In [9]: timeit solve_triangular(L, x)
100 loops, best of 3: 3.45 ms per loop
In [10]: timeit linalg.solve_triangular(L, y, lower=True,
overwrite_b=True)
10 loops, best of 3: 134 ms per loop
[/sourcecode]
Wow! That’s 40x faster. We’re catching two rabbits with one stone here, let’s do the change! Notice that we can just copy [latex] \mathbf{v}[/latex] into the appropriate place in [latex] L[/latex] and then solve in place.
But whoops! When solving the [latex] LL’[/latex] system, we take
advantage of the transpose attribute in linalg.solve_triangular,
which the scikits.learn version does not expose. We could think of a
solution, but here’s a better idea: Shouldn’t there be some way to
directly solve the entire system in one go?
Well, there is. It is an LAPACK function by the name of potrs. If you
are not aware, LAPACK is a Fortran library with solvers for various
types of linear systems and eigenproblems. LAPACK along with BLAS (on
which it is based) pretty much powers all the scientific computation
that happens. BLAS is an API with multiple implementations dating from
1979, while LAPACK dates from 1992. If you ever used Matlab, this is
what was called behind the scenes. SciPy, again, provides a high-level
wrapper around this, the linalg.cho_solve function.
But SciPy also gives us the possibility to import functions directly
from LAPACK, through the use of linalg.lapack.get_lapack_funcs. Let’s
see how the low-level LAPACK function compares to the SciPy wrapper, for
our use case:
[sourcecode language=”python”]
In [11]: x = np.random.randn(1000)
In [12]: y = x.copy()
In [13]: timeit linalg.cho_solve((L, True), x)
1 loops, best of 3: 95.4 ms per loop
In [14]: potrs, = linalg.lapack.get_lapack_funcs((‘potrs’,), (G,))
In [15]: potrs
Out[15]: <fortran object>
In [16]: timeit potrs(L, y)
100 loops, best of 3: 9.49 ms per loop
[/sourcecode]
That’s 10 times faster! So now we found an obvious way to optimize the code:
[sourcecode language=”python”]
def cholesky_omp(X, y, n_nonzero_coefs, eps=None):
min_float = np.finfo(X.dtype).eps
potrs, = get_lapack_funcs((‘potrs’,), (X,))
alpha = np.dot(X.T, y)
residual = y
n_active = 0
idx = []
max_features = X.shape1 if eps is not None else n_nonzero_coefs
L = np.empty((max_features, max_features), dtype=X.dtype)
L[0, 0] = 1.
while 1:
lam = np.abs(np.dot(X.T, residual)).argmax()
if lam < n_active or alpha[lam] ** 2 < min_float:
# atom already selected or inner product too small
warn(“Stopping early”)
break
if n_active > 0:
# Updates the Cholesky decomposition of X’ X
L[n_active, :n_active] = np.dot(X[:, idx].T, X[:, lam]
solve_triangular(L[:n_active, :n_active], L[n_active, :n_active])
d = np.dot(L[n_active, :n_active].T, L[n_active, :n_active])
if 1 - d <= min_float: # selected atoms are dependent
warn(“Stopping early”)
break
L[n_active, n_active] = np.sqrt(1 - d)
idx.append(lam)
# solve LL’x = y in two steps:
gamma, _ = potrs(L[:n_active, :n_active], alpha[idx], lower=True,
overwrite_b=False)
residual = y - np.dot(X[:, idx], gamma)
if eps is not None and np.dot(residual.T, residual) <= eps:
break
elif n_active == max_features:
break
return gamma, idx
[/sourcecode]
Woohoo! But we still lag behind. Now that we delegated the trickiest
parts of the code to fast and reliable solvers, it’s time to use a
profiler and see what the bottleneck is now. Python has excellent tools
for this purpose. What solved the problem in this case was
line_profiler [2]. There is a great article by Olivier Grisel here
2 regarding how to use these profilers. I’m just going to say that
line_profiler‘s output is very helpful, basically printing the time
taken by each line of code next to that line.
Running the profiler on this code, we found that 58% of the time is
spent on line 14, 20.5% on line 21, and 20.5% on line 32, with the rest
being insignificant (potrs takes 0.1%!). The code is clearly dominated
by the matrix multiplications. By running some more timings with
IPython, I found that multiplying such column-wise views of the data as
X[:, idx] is considerably slower then multiplying a contiguous array.
The least-angle regression code in scikits.learn avoids this by
swapping columns towards the front of the array as they are chosen, so
we can replace X[:, idx] with X[:, :n_active]. The nice part is that
if the array is stored in Fortran-contiguous order (ie. column
contiguous order, as opposed to row contiguous order, as in C), swapping
two columns is a very fast operation!. Let’s see some more benchmarks!
[sourcecode language=”python”]
In [17]: X = np.random.randn(5000, 5000)
In [18]: Y = X.copy(‘F’) # fortran-ordered
In [19]: a, b = 1000, 2500
In [20]: swap, = linalg.get_blas_funcs((‘swap’,), (X,))
In [21]: timeit X[:, a], X[:, b] = swap(X[:, a], X[:, b])
100 loops, best of 3: 6.29 ms per loop
In [22]: timeit Y[:, a], Y[:, b] = swap(Y[:, a], Y[:, b])
10000 loops, best of 3: 111 us per loop
[/sourcecode]
We can see that using Fortran-order takes us from the order of miliseconds to the order of microseconds!
Side note: I almost fell into the trap of swapping columns the pythonic
way. That doesn’t work:
[sourcecode language=”python”]
In [23]: X[:, a], X[:, b] = X[:, b], X[:, a]
In [24]: np.testing.assert_array_equal(X[:, a], X[:, b])
In [25]:
[/sourcecode]
However this trick works great for swapping elements of one-dimensional arrays.
Another small optimization that we can do: I found that on my system,
it’s slightly faster to compute the norm using the BLAS function nrm2.
So by putting all of these together, we end up with the final version of
our code:
[sourcecode language=”python”]
def cholesky_omp(X, y, n_nonzero_coefs, eps=None,
overwrite_X=False):
if not overwrite_X:
X = X.copy(‘F’)
else: # even if we are allowed to overwrite, still copy it if bad
order
X = np.asfortranarray(X)
min_float = np.finfo(X.dtype).eps
nrm2, swap = linalg.get_blas_funcs((‘nrm2’, ‘swap’), (X,))
potrs, = get_lapack_funcs((‘potrs’,), (X,))
indices = range(len(Gram)) # keeping track of swapping
alpha = np.dot(X.T, y)
residual = y
n_active = 0
max_features = X.shape1 if eps is not None else n_nonzero_coefs
L = np.empty((max_features, max_features), dtype=X.dtype)
L[0, 0] = 1.
while True:
lam = np.abs(np.dot(X.T, residual)).argmax()
if lam < n_active or alpha[lam] ** 2 < min_float:
# atom already selected or inner product too small
warn(“Stopping early”)
break
if n_active > 0:
# Updates the Cholesky decomposition of X’ X
L[n_active, :n_active] = np.dot(X[:, :n_active].T, X[:, lam])
solve_triangular(L[:n_active, :n_active], L[n_active, :n_active])
v = nrm2(L[n_active, :n_active]) ** 2
if 1 - v <= min_float: # selected atoms are dependent
warn(“Stopping early”)
break
L[n_active, n_active] = np.sqrt(1 - v)
X.T[n_active], X.T[lam] = swap(X.T[n_active], X.T[lam])
alpha[n_active], alpha[lam] = alpha[lam], alpha[n_active]
indices[n_active], indices[lam] = indices[lam], indices[n_active]
n_active += 1
# solves LL’x = y as a composition of two triangular systems
gamma, _ = potrs(L[:n_active, :n_active], alpha[:n_active],
lower=True,
overwrite_b=False)
residual = y - np.dot(X[:, :n_active], gamma)
if eps is not None and nrm2(residual) ** 2 <= eps:
break
elif n_active == max_features:
break
return gamma, indices[:n_active]
[/sourcecode]
Now, the benchmark at [1] indicates victory over least-angle regression! I hope you have enjoyed this short tour. See you next time!
Comments !