[FEA] reduce memory pressure in membership vector computation #5268
Conversation
…n/cuml into fea-membership-vector
…reduce-memory-pressure-apmv
| break; | ||
| case raft::distance::DistanceType::CosineExpanded: | ||
| raft::distance:: | ||
| distance<raft::distance::DistanceType::CosineExpanded, value_t, value_t, value_t, int>( | ||
| handle, X, exemplars_dense.data(), dist.data(), m, n_exemplars, n, true); | ||
| handle, query + batch_offset * n, exemplars_dense.data(), dist.data(), samples_per_batch, n_exemplars, n, true); |
There was a problem hiding this comment.
Very nice! This is a good solution to provide some significant savings in memory.
tarang-jain
force-pushed
the
fea-new-reduce-memory-pressure-apmv
branch
from
91c2d6a to
7a95bfe
Compare
| had ``prediction_data=True`` set. | ||
|
|
||
| batch_size : int, optional, default=0 | ||
| Distance based membership is computed in batches to fit on the device. If not specified, or set to 0, distance based membership is computed at once for all points in the training data. |
There was a problem hiding this comment.
I think this could be worded a little bit better to highlight the purpose of this setting, maybe mention "points in the training data" closer to the beginning? Something like "Lowers memory requirement by computing distance-based membership in smaller batches of points in the training data. Batch size of 0 uses all of the training points, batch size of 1000 computes distances for 1000 points at a time."
| n_prediction_points, | ||
| clusterer.min_samples, | ||
| _metrics_mapping[clusterer.metric], | ||
| <float*> membership_vec_ptr, |
There was a problem hiding this comment.
Just to avoid any potential future issues, we should probably validate that this is actually positive (in Python) before we pass it down.
tarang-jain
added
non-breaking
|
/merge |
Batching is done only while computing the pairwise distance matrix between the dataset and the set of exemplar points.
Closes #4879