cleanlab / cleanlab/cleanlab

Use native code in hot code paths

Open
#297 0 comments 4 reactions 0 assignees View on GitHub
needs triage
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.