quantlab.model.torch_data.prefetch#

quantlab.model.torch_data.prefetch(batches, device)[source]#

Yield batches on device, each readied ahead by a background thread.

A thread takes the next items of batches (a DataLoader, or any iterable of Batch-shaped tuples) in order and moves them to device while the caller works on the current one, keeping at most PREFETCH_DEPTH waiting besides the one it is readying. For a CUDA device it copies each batch into pinned memory and queues the copy to the GPU on a side stream; the caller’s stream waits for that copy before the batch is handed over, so the transfer overlaps the GPU’s work on the previous batch. The batches come out exactly as batches gives them, in its order.

An error raised by batches is raised here, at the batch it stopped at. Closing the iterator early (close(), or an error in the caller’s loop under contextlib.closing) stops and joins the thread.

Parameters:

batchesiterable: The batches, each a Batch or a tuple of its fields.
devicestr or torch.device: Where to hand them over.

Yields:

Batch: The next batch, on device.

Examples

>>> [int(b.where[0][0]) for b in prefetch([batch, batch], "cpu")]
[4, 4]