Use native code in hot code paths
- Dominant language
- Python
- Stars
- 11.7k
- Forks
- 920
- PR merge metrics
- No merged PRs in 30d
Description
We should explore using solutions like Numba or Cython, which can vastly improve performance for hot code paths. There is a downside, though, that introducing tools like these will complicate distribution.
Related to #287; performance benchmarking will help us understand where we should switch to native code. Things like `find_label_issues` could likely be made a _lot_ faster than they are now.
---
As an initial exploration in this direction, I tried using Numba/Cython to speed up [`get_confident_thresholds`](https://github.com/cleanlab/cleanlab/blob/942d97e5b81fb4a887adf518152a9795923496ab/cleanlab/count.py#L1156-L1158). I also made algorithmic improvements. This is what the results look like when evaluated with 50,000 classes and 50,000 examples (pyx size = 18.6 GB), evaluated on an i7 5930K with 32 GB of RAM:
| complexity | implementation | time (sec) | speedup over original | speedup over pure Python/NumPy |
| --- | --- | --: | --: | --: |
| O(mn) | NumPy (original) | 1.7147 | 1x | N/A |
| O(mn) | Numba-slow | 3.0092 | 0.57x | 0.57x |
| O(m + n) | NumPy-fast | 0.1135 | 15x | N/A |
| O(m + n) | Cython-fast | 0.0010 | 1,715x | 114x |
| O(m + n) | Numba-fast | 0.0003 | 5,715x | 378x |
This shows that asymptotic complexity improvements make a big difference (kind of obvious) but also that switching to native code makes a big difference. Below are code snippets / explanations of the different implementations.
**O(mn) NumPy (original)**
The baseline code.
```python
import numpy as np
def get_confident_thresholds(labels, pred_probs):
return np.array([np.mean(pred_probs[:, k][labels == k]) for k in range(pred_probs.shape[1])])
```
**O(mn) Numba-slow**
Just adding a `@numba.jit` to the above:
```python
import numpy as np
import numba
@numba.jit
def get_confident_thresholds(labels, pred_probs):
return np.array([np.mean(pred_probs[:, k][labels == k]) for k in range(pred_probs.shape[1])])
```
Benchmarked by running once to JIT compile the function, and then running again to benchmark it (only timing the second run, so not including JIT compilation time).
**O(m + n) NumPy-fast**
Making an algorithmic improvement to the original NumPy.
```python
import numpy as np
import numba
def get_confident_thresholds(labels, pred_probs):
selected_probs = np.take_along_axis(pred_probs, labels[:,None], axis=1)
sums = np.zeros(pred_probs.shape[1])
counts = np.zeros(pred_probs.shape[1])
for i, p in enumerate(selected_probs):
sums[labels[i]] += p
counts[labels[i]] += 1
return sums / counts
```
**O(m + n) Cython-fast**
Carefully written to minimize allocations and avoid allocations in loops.
```python
import cython
import numpy as np
cimport numpy as np
np.import_array()
@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False) # turn off negative index wrapping for entire function
def fast_thresholds(np.ndarray[np.int64_t, ndim=1] labels, np.ndarray[np.float64_t, ndim=2] pred_probs):
# TODO need to do a single bounds check before doing any indexing
cdef Py_ssize_t num_labels = pred_probs.shape[1] # typed to avoid object allocation
cdef np.ndarray[np.float64_t, ndim=1] sums = np.zeros(num_labels, dtype=np.float64)
cdef np.ndarray[np.int64_t, ndim=1] counts = np.zeros(num_labels, dtype=np.int64)
cdef Py_ssize_t label # typed, and avoid Python object allocation in the loop
for i in range(len(labels)):
label = labels[i]
sums[label] += pred_probs[i, label] # 2-d index to avoid temp object + allocation + decref/deallocation
counts[label] += 1
sums /= counts # in-place divide to avoid allocation
return sums
```
Note to self: investigate Cython [memory views](https://cython.readthedocs.io/en/stable/src/userguide/memoryviews.html).
**O(m + n) Numba-fast**
Porting the above back into Python, and adding a `@numba.jit`.
```python
import numba
import numpy as np
@numba.jit
def numba_fast_thresholds(labels, pred_probs):
num_labels = pred_probs.shape[1]
sums = np.zeros(num_labels, dtype=np.float64)
counts = np.zeros(num_labels, dtype=np.int64)
for i in range(len(labels)):
label = labels[i]
sums[label] += pred_probs[i, label]
counts[label] += 1
sums /= counts
return sums
```
Contributor guide
Research direction
Start with cleanlab/count.py at get_confident_thresholds and review the related find_label_issues path. Reproduce the benchmark comparisons described in the issue, including the algorithmic NumPy version and native-code experiments. Done would require a justified choice of implementation and distribution approach, backed by performance measurements.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100