torchisotonic 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- torchisotonic-0.1.0/.gitignore +14 -0
- torchisotonic-0.1.0/LICENSE +21 -0
- torchisotonic-0.1.0/PKG-INFO +183 -0
- torchisotonic-0.1.0/README.md +158 -0
- torchisotonic-0.1.0/docs/hull_prototype.py +113 -0
- torchisotonic-0.1.0/pyproject.toml +83 -0
- torchisotonic-0.1.0/src/torchisotonic/__init__.py +6 -0
- torchisotonic-0.1.0/src/torchisotonic/_api.py +131 -0
- torchisotonic-0.1.0/src/torchisotonic/_ext.py +52 -0
- torchisotonic-0.1.0/src/torchisotonic/csrc/grad.cuh +422 -0
- torchisotonic-0.1.0/src/torchisotonic/csrc/ops.cpp +50 -0
- torchisotonic-0.1.0/src/torchisotonic/csrc/ops_tiles.cu +13 -0
- torchisotonic-0.1.0/src/torchisotonic/csrc/pava.h +63 -0
- torchisotonic-0.1.0/src/torchisotonic/csrc/tiles.cuh +1920 -0
- torchisotonic-0.1.0/src/torchisotonic/py.typed +0 -0
- torchisotonic-0.1.0/tests/test_cuda.py +218 -0
- torchisotonic-0.1.0/tests/test_grad.py +69 -0
- torchisotonic-0.1.0/tests/test_hull_prototype.py +50 -0
- torchisotonic-0.1.0/tests/test_isotonic.py +163 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Lena Barretta
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: torchisotonic
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Batched isotonic regression for PyTorch tensors, fast on GPU.
|
|
5
|
+
Project-URL: Homepage, https://github.com/LenaBarretta/torchisotonic
|
|
6
|
+
Project-URL: Issues, https://github.com/LenaBarretta/torchisotonic/issues
|
|
7
|
+
Author: Lena Barretta
|
|
8
|
+
License-Expression: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: calibration,cuda,gpu,isotonic-regression,pava,pytorch,torch
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA
|
|
13
|
+
Classifier: Intended Audience :: Developers
|
|
14
|
+
Classifier: Intended Audience :: Science/Research
|
|
15
|
+
Classifier: Programming Language :: C++
|
|
16
|
+
Classifier: Programming Language :: Python :: 3
|
|
17
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
18
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Mathematics
|
|
20
|
+
Classifier: Typing :: Typed
|
|
21
|
+
Requires-Python: >=3.10
|
|
22
|
+
Requires-Dist: ninja
|
|
23
|
+
Requires-Dist: torch>=2.4
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
|
|
26
|
+
# torchisotonic
|
|
27
|
+
|
|
28
|
+
Batched isotonic regression for PyTorch tensors, fast on the GPU.
|
|
29
|
+
|
|
30
|
+
Isotonic regression finds the non-decreasing (or non-increasing) sequence closest to the
|
|
31
|
+
data in weighted least squares. It is what `sklearn.isotonic.isotonic_regression` and
|
|
32
|
+
`scipy.optimize.isotonic_regression` compute for one sequence; torchisotonic computes it for
|
|
33
|
+
every row of a tensor at once, on the CPU or with CUDA kernels on the GPU, from a few
|
|
34
|
+
elements per row to a hundred million.
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
import torch
|
|
38
|
+
import torchisotonic
|
|
39
|
+
|
|
40
|
+
y = torch.randn(1000, 5000, device="cuda")
|
|
41
|
+
fit = torchisotonic.isotonic(y) # every row, non-decreasing
|
|
42
|
+
|
|
43
|
+
w = torch.rand(1000, 5000, device="cuda") + 0.1
|
|
44
|
+
fit = torchisotonic.isotonic(y, w, increasing=False, y_min=0.0, y_max=1.0)
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
## What it is for
|
|
48
|
+
|
|
49
|
+
Whenever a monotone fit has to run often, on long data, for many curves at once, or inside a
|
|
50
|
+
model during training.
|
|
51
|
+
|
|
52
|
+
**Calibrating a classifier.** A model's scores rank well but are not probabilities; isotonic
|
|
53
|
+
regression maps each score to the rate of positives observed at that score, as scikit-learn's
|
|
54
|
+
`IsotonicRegression` and `CalibratedClassifierCV(method="isotonic")` do. A curve over 6.5 million
|
|
55
|
+
held-out scores takes 42 ms on a T4, against 2 s for scikit-learn. These two functions give
|
|
56
|
+
scikit-learn's curve and predictions, equal scores pooled into their mean:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
def fit_calibration(scores, labels):
|
|
60
|
+
"""The calibration curve, as scikit-learn's IsotonicRegression fits it."""
|
|
61
|
+
s, order = torch.sort(scores)
|
|
62
|
+
xs, inverse, counts = torch.unique_consecutive(s, return_inverse=True, return_counts=True)
|
|
63
|
+
sums = torch.zeros(len(xs), dtype=torch.float64, device=scores.device)
|
|
64
|
+
sums.index_add_(0, inverse, labels[order].double())
|
|
65
|
+
return xs, torchisotonic.isotonic(sums / counts, counts.double())
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def calibrate(new_scores, xs, fit):
|
|
69
|
+
"""Calibrated probabilities: linear between the curve's points, constant beyond them."""
|
|
70
|
+
i = torch.searchsorted(xs, new_scores).clamp(1, len(xs) - 1)
|
|
71
|
+
t = ((new_scores - xs[i - 1]) / (xs[i] - xs[i - 1])).clamp(0, 1)
|
|
72
|
+
return fit[i - 1] + t * (fit[i] - fit[i - 1])
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
[`examples/ctr_calibration`](examples/ctr_calibration) calibrates a click model on 45.8
|
|
76
|
+
million Criteo impressions this way, with one curve per segment for 280 segments in one call.
|
|
77
|
+
|
|
78
|
+
**Many curves at once.** Every row of the input is its own problem, so a thousand curves cost
|
|
79
|
+
little more than one: per segment, per threshold, or per bootstrap sample. For instance, a
|
|
80
|
+
band of uncertainty around a calibration curve, from a Bayesian bootstrap of the labels `y`
|
|
81
|
+
ordered by score:
|
|
82
|
+
|
|
83
|
+
```python
|
|
84
|
+
w = torch.empty(1000, len(y), device=y.device).exponential_() # one resample per row
|
|
85
|
+
curves = torchisotonic.isotonic(y.expand(1000, -1), w)
|
|
86
|
+
low, high = curves.quantile(0.05, dim=0), curves.quantile(0.95, dim=0)
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
**Inside a model.** The fit is differentiable in `y` and in the weights, so it can enforce
|
|
90
|
+
monotonicity in a network and train through it: predicted quantiles that never cross,
|
|
91
|
+
dose-response or survival curves that must not go the wrong way, or the projection step of
|
|
92
|
+
differentiable sorting and ranking.
|
|
93
|
+
|
|
94
|
+
```python
|
|
95
|
+
q = model(x) # [batch, levels]: quantiles at increasing levels
|
|
96
|
+
q = torchisotonic.isotonic(q) # never crossing, with gradients
|
|
97
|
+
loss = pinball(q, target, levels) # any quantile loss
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
## Semantics
|
|
101
|
+
|
|
102
|
+
The semantics are scikit-learn's: weights must be strictly positive (a weight of the shape of
|
|
103
|
+
`y`, or one that broadcasts to it), the result is clipped to `[y_min, y_max]` if they are
|
|
104
|
+
given, `dim` picks the axis, and a row holding a NaN comes back as NaN. Float input keeps its
|
|
105
|
+
dtype (float16 and bfloat16 are computed in float32), integer input becomes the default
|
|
106
|
+
float dtype. The fit is differentiable in `y` and in the weights.
|
|
107
|
+
|
|
108
|
+
On the GPU, float32 is accurate to the format's rounding by default, and the fit never steps
|
|
109
|
+
down. `exact=False` makes it about 1.5–2 times faster on short and medium rows, at the cost of
|
|
110
|
+
accuracy on hard inputs (strongly varying weights, increments near float32's resolution): errors
|
|
111
|
+
of a few parts in a thousand, and steps down of up to about 1e-3 of the data's scale; float64
|
|
112
|
+
is unaffected.
|
|
113
|
+
|
|
114
|
+
Rows of unequal length go in one batch padded at the end with a value above each row's maximum
|
|
115
|
+
(below its minimum for `increasing=False`): the padding never pools with the data, so the fit
|
|
116
|
+
of each row is exactly its own. Rows of very different lengths fit faster laid end to end in a single row,
|
|
117
|
+
as [`examples/ctr_calibration`](examples/ctr_calibration) does.
|
|
118
|
+
|
|
119
|
+
## Installation
|
|
120
|
+
|
|
121
|
+
```bash
|
|
122
|
+
pip install torchisotonic
|
|
123
|
+
```
|
|
124
|
+
|
|
125
|
+
The C++/CUDA code is compiled the first time it is used (about a minute), which needs `ninja`,
|
|
126
|
+
a C++ compiler and, for the GPU, the CUDA toolkit. **Status: alpha**; the API may still change.
|
|
127
|
+
|
|
128
|
+
## Performance
|
|
129
|
+
|
|
130
|
+
On an NVIDIA T4 (Kaggle), float32, unweighted, decreasing fits of data noisy around a trend,
|
|
131
|
+
in milliseconds; best of ten calls on a warmed-up GPU, data already on the device for the GPU
|
|
132
|
+
columns:
|
|
133
|
+
|
|
134
|
+
| batch × n | torchisotonic (cuda) | `exact=False` | torchsort (cuda) | torchisotonic (cpu) | torchsort (cpu) | scipy, a loop over rows |
|
|
135
|
+
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
|
|
136
|
+
| 1,000,000 × 16 | **2.16** | 1.21 | 72.0 | 407 | 677 | 14,363 |
|
|
137
|
+
| 100,000 × 128 | **1.71** | 1.30 | 136 | 280 | 492 | 1,613 |
|
|
138
|
+
| 10,000 × 1,000 | **2.29** | 1.22 | 99.2 | 160 | 354 | 318 |
|
|
139
|
+
| 1,000 × 10,000 | **1.72** | 1.12 | 197 | 160 | 358 | 178 |
|
|
140
|
+
| 100 × 100,000 | **1.82** | 1.23 | 396 | 158 | 352 | 166 |
|
|
141
|
+
| 10 × 1,000,000 | **1.95** | 1.35 | 1,656 | 157 | 357 | 170 |
|
|
142
|
+
| 1 × 10,000,000 | **2.07** | 1.47 | 5,493 | 236 | 355 | 176 |
|
|
143
|
+
| 100 × 1,000 | **0.09** | 0.07 | 3.0 | 1.7 | 3.5 | 2.9 |
|
|
144
|
+
| 10 × 10,000 | **0.22** | 0.20 | 14.0 | 1.7 | 3.5 | 1.6 |
|
|
145
|
+
| 1 × 100,000 | **0.26** | 0.25 | 51.7 | 1.8 | 3.5 | 1.6 |
|
|
146
|
+
| 1 × 1,000,000 | **0.51** | 0.44 | 550 | 17.9 | 35.0 | 15.5 |
|
|
147
|
+
| 1 × 100,000,000 | **16.9** | 10.8 | — | 3,552 | — | — |
|
|
148
|
+
|
|
149
|
+
torchsort solves each row with one GPU thread, so a long row is slow; it was not run on 10⁸
|
|
150
|
+
elements. The CPU columns are on Kaggle's CPU, so they say how the libraries compare, not how
|
|
151
|
+
fast a good CPU is. On an L4, an A100 and an H100 one row of 10⁷ elements takes 1.2, 0.8 and
|
|
152
|
+
0.5 ms, and 10⁸ elements 8.2, 4.3 and 2.7 ms. `bench/bench.py` reproduces the table.
|
|
153
|
+
|
|
154
|
+
## How it works
|
|
155
|
+
|
|
156
|
+
The fit at each element is the slope of the lower convex hull of the cumulative sums of the
|
|
157
|
+
data (the "greatest convex minorant"), so the GPU code builds that hull in parallel:
|
|
158
|
+
|
|
159
|
+
1. Every lane of a warp builds the hull of 32 neighbouring points with the monotone chain,
|
|
160
|
+
its stack kept as a bit mask; then the lane hulls of a chunk of 1024 to 4096 elements merge
|
|
161
|
+
pairwise in shared memory, each merge walking to its bridge and clearing the bits under it.
|
|
162
|
+
Short rows are packed several per warp. A row that fits in one chunk is finished there.
|
|
163
|
+
2. Longer rows merge their chunks' hulls pairwise, in rounds, as lists of segments of the
|
|
164
|
+
chunks' hulls, so that a merge costs segments rather than vertices.
|
|
165
|
+
3. Every element takes the slope of the hull edge above it.
|
|
166
|
+
|
|
167
|
+
[`docs/hull_prototype.py`](docs/hull_prototype.py) is a plain-numpy reference of the parallel
|
|
168
|
+
algorithm; [`src/torchisotonic/csrc/tiles.cuh`](src/torchisotonic/csrc/tiles.cuh)
|
|
169
|
+
describes the kernels. On the CPU, rows are solved with PAVA, in parallel.
|
|
170
|
+
|
|
171
|
+
## Development
|
|
172
|
+
|
|
173
|
+
```bash
|
|
174
|
+
uv sync
|
|
175
|
+
uv run pytest
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
The tests compare against scipy and scikit-learn (with hypothesis) and, on a machine with a
|
|
179
|
+
GPU, the CUDA operator against the CPU one; without a GPU the CUDA tests are skipped.
|
|
180
|
+
|
|
181
|
+
## License
|
|
182
|
+
|
|
183
|
+
MIT
|
|
@@ -0,0 +1,158 @@
|
|
|
1
|
+
# torchisotonic
|
|
2
|
+
|
|
3
|
+
Batched isotonic regression for PyTorch tensors, fast on the GPU.
|
|
4
|
+
|
|
5
|
+
Isotonic regression finds the non-decreasing (or non-increasing) sequence closest to the
|
|
6
|
+
data in weighted least squares. It is what `sklearn.isotonic.isotonic_regression` and
|
|
7
|
+
`scipy.optimize.isotonic_regression` compute for one sequence; torchisotonic computes it for
|
|
8
|
+
every row of a tensor at once, on the CPU or with CUDA kernels on the GPU, from a few
|
|
9
|
+
elements per row to a hundred million.
|
|
10
|
+
|
|
11
|
+
```python
|
|
12
|
+
import torch
|
|
13
|
+
import torchisotonic
|
|
14
|
+
|
|
15
|
+
y = torch.randn(1000, 5000, device="cuda")
|
|
16
|
+
fit = torchisotonic.isotonic(y) # every row, non-decreasing
|
|
17
|
+
|
|
18
|
+
w = torch.rand(1000, 5000, device="cuda") + 0.1
|
|
19
|
+
fit = torchisotonic.isotonic(y, w, increasing=False, y_min=0.0, y_max=1.0)
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
## What it is for
|
|
23
|
+
|
|
24
|
+
Whenever a monotone fit has to run often, on long data, for many curves at once, or inside a
|
|
25
|
+
model during training.
|
|
26
|
+
|
|
27
|
+
**Calibrating a classifier.** A model's scores rank well but are not probabilities; isotonic
|
|
28
|
+
regression maps each score to the rate of positives observed at that score, as scikit-learn's
|
|
29
|
+
`IsotonicRegression` and `CalibratedClassifierCV(method="isotonic")` do. A curve over 6.5 million
|
|
30
|
+
held-out scores takes 42 ms on a T4, against 2 s for scikit-learn. These two functions give
|
|
31
|
+
scikit-learn's curve and predictions, equal scores pooled into their mean:
|
|
32
|
+
|
|
33
|
+
```python
|
|
34
|
+
def fit_calibration(scores, labels):
|
|
35
|
+
"""The calibration curve, as scikit-learn's IsotonicRegression fits it."""
|
|
36
|
+
s, order = torch.sort(scores)
|
|
37
|
+
xs, inverse, counts = torch.unique_consecutive(s, return_inverse=True, return_counts=True)
|
|
38
|
+
sums = torch.zeros(len(xs), dtype=torch.float64, device=scores.device)
|
|
39
|
+
sums.index_add_(0, inverse, labels[order].double())
|
|
40
|
+
return xs, torchisotonic.isotonic(sums / counts, counts.double())
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def calibrate(new_scores, xs, fit):
|
|
44
|
+
"""Calibrated probabilities: linear between the curve's points, constant beyond them."""
|
|
45
|
+
i = torch.searchsorted(xs, new_scores).clamp(1, len(xs) - 1)
|
|
46
|
+
t = ((new_scores - xs[i - 1]) / (xs[i] - xs[i - 1])).clamp(0, 1)
|
|
47
|
+
return fit[i - 1] + t * (fit[i] - fit[i - 1])
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
[`examples/ctr_calibration`](examples/ctr_calibration) calibrates a click model on 45.8
|
|
51
|
+
million Criteo impressions this way, with one curve per segment for 280 segments in one call.
|
|
52
|
+
|
|
53
|
+
**Many curves at once.** Every row of the input is its own problem, so a thousand curves cost
|
|
54
|
+
little more than one: per segment, per threshold, or per bootstrap sample. For instance, a
|
|
55
|
+
band of uncertainty around a calibration curve, from a Bayesian bootstrap of the labels `y`
|
|
56
|
+
ordered by score:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
w = torch.empty(1000, len(y), device=y.device).exponential_() # one resample per row
|
|
60
|
+
curves = torchisotonic.isotonic(y.expand(1000, -1), w)
|
|
61
|
+
low, high = curves.quantile(0.05, dim=0), curves.quantile(0.95, dim=0)
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
**Inside a model.** The fit is differentiable in `y` and in the weights, so it can enforce
|
|
65
|
+
monotonicity in a network and train through it: predicted quantiles that never cross,
|
|
66
|
+
dose-response or survival curves that must not go the wrong way, or the projection step of
|
|
67
|
+
differentiable sorting and ranking.
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
q = model(x) # [batch, levels]: quantiles at increasing levels
|
|
71
|
+
q = torchisotonic.isotonic(q) # never crossing, with gradients
|
|
72
|
+
loss = pinball(q, target, levels) # any quantile loss
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
## Semantics
|
|
76
|
+
|
|
77
|
+
The semantics are scikit-learn's: weights must be strictly positive (a weight of the shape of
|
|
78
|
+
`y`, or one that broadcasts to it), the result is clipped to `[y_min, y_max]` if they are
|
|
79
|
+
given, `dim` picks the axis, and a row holding a NaN comes back as NaN. Float input keeps its
|
|
80
|
+
dtype (float16 and bfloat16 are computed in float32), integer input becomes the default
|
|
81
|
+
float dtype. The fit is differentiable in `y` and in the weights.
|
|
82
|
+
|
|
83
|
+
On the GPU, float32 is accurate to the format's rounding by default, and the fit never steps
|
|
84
|
+
down. `exact=False` makes it about 1.5–2 times faster on short and medium rows, at the cost of
|
|
85
|
+
accuracy on hard inputs (strongly varying weights, increments near float32's resolution): errors
|
|
86
|
+
of a few parts in a thousand, and steps down of up to about 1e-3 of the data's scale; float64
|
|
87
|
+
is unaffected.
|
|
88
|
+
|
|
89
|
+
Rows of unequal length go in one batch padded at the end with a value above each row's maximum
|
|
90
|
+
(below its minimum for `increasing=False`): the padding never pools with the data, so the fit
|
|
91
|
+
of each row is exactly its own. Rows of very different lengths fit faster laid end to end in a single row,
|
|
92
|
+
as [`examples/ctr_calibration`](examples/ctr_calibration) does.
|
|
93
|
+
|
|
94
|
+
## Installation
|
|
95
|
+
|
|
96
|
+
```bash
|
|
97
|
+
pip install torchisotonic
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
The C++/CUDA code is compiled the first time it is used (about a minute), which needs `ninja`,
|
|
101
|
+
a C++ compiler and, for the GPU, the CUDA toolkit. **Status: alpha**; the API may still change.
|
|
102
|
+
|
|
103
|
+
## Performance
|
|
104
|
+
|
|
105
|
+
On an NVIDIA T4 (Kaggle), float32, unweighted, decreasing fits of data noisy around a trend,
|
|
106
|
+
in milliseconds; best of ten calls on a warmed-up GPU, data already on the device for the GPU
|
|
107
|
+
columns:
|
|
108
|
+
|
|
109
|
+
| batch × n | torchisotonic (cuda) | `exact=False` | torchsort (cuda) | torchisotonic (cpu) | torchsort (cpu) | scipy, a loop over rows |
|
|
110
|
+
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
|
|
111
|
+
| 1,000,000 × 16 | **2.16** | 1.21 | 72.0 | 407 | 677 | 14,363 |
|
|
112
|
+
| 100,000 × 128 | **1.71** | 1.30 | 136 | 280 | 492 | 1,613 |
|
|
113
|
+
| 10,000 × 1,000 | **2.29** | 1.22 | 99.2 | 160 | 354 | 318 |
|
|
114
|
+
| 1,000 × 10,000 | **1.72** | 1.12 | 197 | 160 | 358 | 178 |
|
|
115
|
+
| 100 × 100,000 | **1.82** | 1.23 | 396 | 158 | 352 | 166 |
|
|
116
|
+
| 10 × 1,000,000 | **1.95** | 1.35 | 1,656 | 157 | 357 | 170 |
|
|
117
|
+
| 1 × 10,000,000 | **2.07** | 1.47 | 5,493 | 236 | 355 | 176 |
|
|
118
|
+
| 100 × 1,000 | **0.09** | 0.07 | 3.0 | 1.7 | 3.5 | 2.9 |
|
|
119
|
+
| 10 × 10,000 | **0.22** | 0.20 | 14.0 | 1.7 | 3.5 | 1.6 |
|
|
120
|
+
| 1 × 100,000 | **0.26** | 0.25 | 51.7 | 1.8 | 3.5 | 1.6 |
|
|
121
|
+
| 1 × 1,000,000 | **0.51** | 0.44 | 550 | 17.9 | 35.0 | 15.5 |
|
|
122
|
+
| 1 × 100,000,000 | **16.9** | 10.8 | — | 3,552 | — | — |
|
|
123
|
+
|
|
124
|
+
torchsort solves each row with one GPU thread, so a long row is slow; it was not run on 10⁸
|
|
125
|
+
elements. The CPU columns are on Kaggle's CPU, so they say how the libraries compare, not how
|
|
126
|
+
fast a good CPU is. On an L4, an A100 and an H100 one row of 10⁷ elements takes 1.2, 0.8 and
|
|
127
|
+
0.5 ms, and 10⁸ elements 8.2, 4.3 and 2.7 ms. `bench/bench.py` reproduces the table.
|
|
128
|
+
|
|
129
|
+
## How it works
|
|
130
|
+
|
|
131
|
+
The fit at each element is the slope of the lower convex hull of the cumulative sums of the
|
|
132
|
+
data (the "greatest convex minorant"), so the GPU code builds that hull in parallel:
|
|
133
|
+
|
|
134
|
+
1. Every lane of a warp builds the hull of 32 neighbouring points with the monotone chain,
|
|
135
|
+
its stack kept as a bit mask; then the lane hulls of a chunk of 1024 to 4096 elements merge
|
|
136
|
+
pairwise in shared memory, each merge walking to its bridge and clearing the bits under it.
|
|
137
|
+
Short rows are packed several per warp. A row that fits in one chunk is finished there.
|
|
138
|
+
2. Longer rows merge their chunks' hulls pairwise, in rounds, as lists of segments of the
|
|
139
|
+
chunks' hulls, so that a merge costs segments rather than vertices.
|
|
140
|
+
3. Every element takes the slope of the hull edge above it.
|
|
141
|
+
|
|
142
|
+
[`docs/hull_prototype.py`](docs/hull_prototype.py) is a plain-numpy reference of the parallel
|
|
143
|
+
algorithm; [`src/torchisotonic/csrc/tiles.cuh`](src/torchisotonic/csrc/tiles.cuh)
|
|
144
|
+
describes the kernels. On the CPU, rows are solved with PAVA, in parallel.
|
|
145
|
+
|
|
146
|
+
## Development
|
|
147
|
+
|
|
148
|
+
```bash
|
|
149
|
+
uv sync
|
|
150
|
+
uv run pytest
|
|
151
|
+
```
|
|
152
|
+
|
|
153
|
+
The tests compare against scipy and scikit-learn (with hypothesis) and, on a machine with a
|
|
154
|
+
GPU, the CUDA operator against the CPU one; without a GPU the CUDA tests are skipped.
|
|
155
|
+
|
|
156
|
+
## License
|
|
157
|
+
|
|
158
|
+
MIT
|
|
@@ -0,0 +1,113 @@
|
|
|
1
|
+
"""Reference prototype of the parallel algorithm for one large problem, in plain numpy.
|
|
2
|
+
|
|
3
|
+
Isotonic regression is the slope of the greatest convex minorant of the cumulative sum
|
|
4
|
+
diagram (Barlow et al.; see also Kearsley, J. Res. NIST 111(2), 2006): with points
|
|
5
|
+
P_k = (W_k, S_k), W_k = w_0 + ... + w_{k-1} and S_k = w_0 y_0 + ... + w_{k-1} y_{k-1}, the fit
|
|
6
|
+
at element i is the slope of the lower convex hull of P_0..P_n over [W_i, W_{i+1}].
|
|
7
|
+
|
|
8
|
+
The GPU version does this in parallel:
|
|
9
|
+
|
|
10
|
+
1. Cut the problem into chunks of C elements; solve each with PAVA, independently. A chunk's
|
|
11
|
+
blocks are the edges of its own lower hull, so its block boundaries are hull vertices.
|
|
12
|
+
2. Merge neighbouring hulls in log2(#chunks) rounds. Two neighbouring lower hulls share their
|
|
13
|
+
boundary point. If the hull is already convex there, they just concatenate; otherwise
|
|
14
|
+
there is one bridge, a common lower tangent from vertex p of the left hull to vertex q of
|
|
15
|
+
the right one. The vertices strictly between p and q drop out, and the elements between
|
|
16
|
+
them form one block.
|
|
17
|
+
3. Every element takes the slope of the final hull edge above it.
|
|
18
|
+
|
|
19
|
+
The bridge is found the way the GPU does it, all candidates at once: every left vertex p
|
|
20
|
+
finds its own lower tangent to the right hull by binary search (slopes from p to the right
|
|
21
|
+
hull's vertices fall, then rise), and the bridge is the p whose tangent also supports the
|
|
22
|
+
left hull at p (or, with rounding, comes closest to it).
|
|
23
|
+
|
|
24
|
+
This file only checks the logic (tests/test_hull_prototype.py runs it against scipy). The
|
|
25
|
+
coordinates here are plain float64 prefix sums; the CUDA kernel works in local coordinates.
|
|
26
|
+
"""
|
|
27
|
+
|
|
28
|
+
from __future__ import annotations
|
|
29
|
+
|
|
30
|
+
import itertools
|
|
31
|
+
|
|
32
|
+
import numpy as np
|
|
33
|
+
import numpy.typing as npt
|
|
34
|
+
|
|
35
|
+
Array = npt.NDArray[np.float64]
|
|
36
|
+
Hull = list[int] # indices k of the points P_k on a lower hull, increasing
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def chunk_hull(w: Array, y: Array, lo: int, hi: int) -> Hull:
|
|
40
|
+
"""PAVA on elements lo..hi-1; returns its hull as point indices lo..hi."""
|
|
41
|
+
starts: list[int] = []
|
|
42
|
+
weights: list[float] = []
|
|
43
|
+
totals: list[float] = []
|
|
44
|
+
for i in range(lo, hi):
|
|
45
|
+
starts.append(i)
|
|
46
|
+
weights.append(float(w[i]))
|
|
47
|
+
totals.append(float(w[i] * y[i]))
|
|
48
|
+
while len(starts) > 1 and totals[-2] * weights[-1] >= totals[-1] * weights[-2]:
|
|
49
|
+
top_weight, top_total = weights.pop(), totals.pop()
|
|
50
|
+
weights[-1] += top_weight
|
|
51
|
+
totals[-1] += top_total
|
|
52
|
+
starts.pop()
|
|
53
|
+
return [*starts, hi]
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def slope(xs: Array, ys: Array, a: int, b: int) -> float:
|
|
57
|
+
return float((ys[b] - ys[a]) / (xs[b] - xs[a]))
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def tangent(xs: Array, ys: Array, p: int, right: Hull) -> int:
|
|
61
|
+
"""Position in ``right`` of the vertex where the lower tangent from point p touches."""
|
|
62
|
+
lo, hi = 0, len(right) - 1
|
|
63
|
+
while lo < hi: # slopes from p fall then rise along the right hull: find the minimum
|
|
64
|
+
mid = (lo + hi) // 2
|
|
65
|
+
if slope(xs, ys, p, right[mid + 1]) < slope(xs, ys, p, right[mid]):
|
|
66
|
+
lo = mid + 1
|
|
67
|
+
else:
|
|
68
|
+
hi = mid
|
|
69
|
+
return lo
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def merge(xs: Array, ys: Array, left: Hull, right: Hull) -> Hull:
|
|
73
|
+
"""Lower hull of two neighbouring lower hulls (left[-1] == right[0])."""
|
|
74
|
+
assert left[-1] == right[0]
|
|
75
|
+
if len(left) == 1 or len(right) == 1:
|
|
76
|
+
return left + right[1:]
|
|
77
|
+
if slope(xs, ys, left[-2], left[-1]) <= slope(xs, ys, right[0], right[1]):
|
|
78
|
+
return left + right[1:] # already convex at the shared point
|
|
79
|
+
# Every left vertex but the shared one is a candidate; on the GPU, one thread each. In
|
|
80
|
+
# exact arithmetic exactly one has a tangent that also supports the left hull (violation
|
|
81
|
+
# 0). Rounding can make the tests contradict each other on (nearly) collinear points, so
|
|
82
|
+
# take the candidate that violates support least: an argmin, i.e. one GPU reduction.
|
|
83
|
+
best, best_violation, best_b = 0, float("inf"), 0
|
|
84
|
+
for a in range(len(left) - 1):
|
|
85
|
+
p = left[a]
|
|
86
|
+
b = tangent(xs, ys, p, right[1:]) + 1
|
|
87
|
+
s = slope(xs, ys, p, right[b])
|
|
88
|
+
before = 0.0 if a == 0 else slope(xs, ys, left[a - 1], p) - s
|
|
89
|
+
after = s - slope(xs, ys, p, left[a + 1])
|
|
90
|
+
violation = max(before, after, 0.0)
|
|
91
|
+
if violation < best_violation:
|
|
92
|
+
best, best_violation, best_b = a, violation, b
|
|
93
|
+
return left[: best + 1] + right[best_b:]
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def isotonic(y: Array, w: Array | None = None, chunk: int = 4) -> Array:
|
|
97
|
+
"""Increasing isotonic regression of y, computed the parallel way."""
|
|
98
|
+
n = len(y)
|
|
99
|
+
if n == 0:
|
|
100
|
+
return y.astype(np.float64)
|
|
101
|
+
w = np.ones(n) if w is None else np.asarray(w, dtype=np.float64)
|
|
102
|
+
y = np.asarray(y, dtype=np.float64)
|
|
103
|
+
xs = np.concatenate([[0.0], np.cumsum(w)])
|
|
104
|
+
ys = np.concatenate([[0.0], np.cumsum(w * y)])
|
|
105
|
+
hulls = [chunk_hull(w, y, lo, min(lo + chunk, n)) for lo in range(0, n, chunk)]
|
|
106
|
+
while len(hulls) > 1: # one round of pairwise merges per level of the tree
|
|
107
|
+
merged = [merge(xs, ys, hulls[i], hulls[i + 1]) for i in range(0, len(hulls) - 1, 2)]
|
|
108
|
+
hulls = merged + ([hulls[-1]] if len(hulls) % 2 else [])
|
|
109
|
+
vertices = hulls[0]
|
|
110
|
+
out = np.empty(n)
|
|
111
|
+
for a, b in itertools.pairwise(vertices):
|
|
112
|
+
out[a:b] = slope(xs, ys, a, b)
|
|
113
|
+
return out
|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling>=1.24"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "torchisotonic"
|
|
7
|
+
dynamic = ["version"]
|
|
8
|
+
description = "Batched isotonic regression for PyTorch tensors, fast on GPU."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
license-files = ["LICENSE"]
|
|
12
|
+
authors = [{ name = "Lena Barretta" }]
|
|
13
|
+
requires-python = ">=3.10"
|
|
14
|
+
# The C++/CUDA operators are compiled on first use, which needs ninja and a C++ compiler.
|
|
15
|
+
dependencies = ["torch>=2.4", "ninja"]
|
|
16
|
+
keywords = ["pytorch", "torch", "isotonic-regression", "calibration", "pava", "cuda", "gpu"]
|
|
17
|
+
classifiers = [
|
|
18
|
+
"Development Status :: 3 - Alpha",
|
|
19
|
+
"Environment :: GPU :: NVIDIA CUDA",
|
|
20
|
+
"Intended Audience :: Developers",
|
|
21
|
+
"Intended Audience :: Science/Research",
|
|
22
|
+
"Programming Language :: Python :: 3",
|
|
23
|
+
"Programming Language :: Python :: 3 :: Only",
|
|
24
|
+
"Programming Language :: C++",
|
|
25
|
+
"Topic :: Scientific/Engineering :: Mathematics",
|
|
26
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
27
|
+
"Typing :: Typed",
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
[project.urls]
|
|
31
|
+
Homepage = "https://github.com/LenaBarretta/torchisotonic"
|
|
32
|
+
Issues = "https://github.com/LenaBarretta/torchisotonic/issues"
|
|
33
|
+
|
|
34
|
+
[dependency-groups]
|
|
35
|
+
dev = [
|
|
36
|
+
"pytest>=8",
|
|
37
|
+
"pytest-cov>=5",
|
|
38
|
+
"ruff>=0.6",
|
|
39
|
+
"mypy>=1.11",
|
|
40
|
+
"hypothesis>=6",
|
|
41
|
+
"numpy>=1.24",
|
|
42
|
+
"scipy>=1.12",
|
|
43
|
+
"scikit-learn>=1.4",
|
|
44
|
+
]
|
|
45
|
+
|
|
46
|
+
# Develop and test against the small CPU-only torch build.
|
|
47
|
+
[tool.uv.sources]
|
|
48
|
+
torch = [{ index = "pytorch-cpu" }]
|
|
49
|
+
|
|
50
|
+
[[tool.uv.index]]
|
|
51
|
+
name = "pytorch-cpu"
|
|
52
|
+
url = "https://download.pytorch.org/whl/cpu"
|
|
53
|
+
explicit = true
|
|
54
|
+
|
|
55
|
+
[tool.hatch.build.targets.sdist]
|
|
56
|
+
# The sources, the tests (with docs/hull_prototype.py, which one of them checks), README and
|
|
57
|
+
# licence. notes/ and IDE settings are kept out of git through .git/info/exclude, which hatch
|
|
58
|
+
# does not read, so they are listed here too.
|
|
59
|
+
exclude = ["bench", "examples", ".github", "notes", ".idea", ".vscode", "uv.lock"]
|
|
60
|
+
|
|
61
|
+
[tool.hatch.version]
|
|
62
|
+
path = "src/torchisotonic/__init__.py"
|
|
63
|
+
|
|
64
|
+
[tool.pytest.ini_options]
|
|
65
|
+
addopts = "-q --cov=torchisotonic --cov-report=term-missing --cov-fail-under=100"
|
|
66
|
+
testpaths = ["tests"]
|
|
67
|
+
|
|
68
|
+
[tool.ruff]
|
|
69
|
+
line-length = 100
|
|
70
|
+
target-version = "py310"
|
|
71
|
+
|
|
72
|
+
[tool.ruff.lint]
|
|
73
|
+
select = ["E", "F", "W", "I", "B", "UP", "SIM", "RUF"]
|
|
74
|
+
|
|
75
|
+
[tool.mypy]
|
|
76
|
+
strict = true
|
|
77
|
+
files = ["src", "tests", "docs"]
|
|
78
|
+
# tests/test_hull_prototype.py imports docs/hull_prototype.py by putting docs/ on sys.path.
|
|
79
|
+
mypy_path = "docs"
|
|
80
|
+
|
|
81
|
+
[[tool.mypy.overrides]]
|
|
82
|
+
module = ["scipy", "scipy.*", "sklearn", "sklearn.*"]
|
|
83
|
+
ignore_missing_imports = true
|