feat: enable composable and customizable sampler in PyTorch data loader #1900
Conversation
89faedd to
48eb133
Compare
| @@ -184,3 +186,130 @@ def reservoir_sampling(stream: Iterable[T], k: int) -> list[T]: | |||
| samples = [i.item for i in heap] | |||
| del heap | |||
| return samples | |||
|
|
|||
|
|
|||
| class Sampler(ABC): | |||
There was a problem hiding this comment.
Right now the implementations just scan in order, they don't randomize (which almost makes them not really meet the definition of "sampling".) Users could do shuffling / reservoir sampling on the batches, but it would much more efficient to do it on fragment_ids and batch indices. Do you have any plans to integrate that with this API?
There was a problem hiding this comment.
Reservoir shuffling is I/O friendly with sequential read (which is NFS friendly), while yielding random batch in uniformity distribution. The time complexity is O(k + log(n/k)) where n is the # of batches, and k is small, with k memory foot prints, and amortizes the file path lookup and metadata overhead cross the scan. Within Lance itself, it is more performant to run read_batch than take. In many cases, Reservoir shuffling can provide pretty decent performance. Need more performance numbers for sure.
I might need to get another PR out to put np.random.select(fragments) and reservior_shuffle(batches) tho. This one established the APIs.
There was a problem hiding this comment.
That being said, the reservior sampling need to be change to
def reservoir_sampling(stream: Iterable[T], k: int, rank: int, world_sizae: int) -> list[T]:
rng = np.random.default_rng()
heap = []
for idx, item in enumerate(stream):
entry = PrioritizedItem(rng.integers(0, k * 2), item)
if len(heap) < k:
heappush(heap, entry)
else:
vic = heappushpop(heap, entry)
if idx % world_size == rank: ## <<<<< CHANGE TO YIELD HERE
yield vic
del vic
if idx % 10240 == 0:
logging.info("Force Python GC")
gc.collect()
samples = [i.item for i in heap]
del heap
return samplesRun this with n=1M, k=8, world_size=1
python/python/lance/sampler.py
Outdated
| Each rank / process will process a subset of the batches. | ||
| """ | ||
|
|
||
| def __init__(self, rank: int, world_size: int): |
There was a problem hiding this comment.
nit: add a from_torch method
There was a problem hiding this comment.
the one read from torch distributed?
python/python/lance/sampler.py
Outdated
| Each rank / process will process a subset of the fragments. | ||
| """ | ||
|
|
||
| def __init__(self, rank: int, world_size: int): |
ruffmade a bunch of format changes