CudaDataLoader

class composer.trainer.devices.device_gpu.CudaDataLoader(dataloader: DataLoader, prefetch_in_cuda_stream: bool, device: torch.device, prefetch_fn: Optional[TPrefetchFn] = None)[source]

Bases: composer.datasets.dataloader.WrappedDataLoader

Wraps DataLoader and moves samples onto the specified device as they are used.

Parameters

dataloader (DataLoader) – The dataloader to wrap.
prefetch_in_cuda_stream (bool) – True to asyncrhonously prefetch samples with a CUDA stream during dataloading and False otherwise.
device (torch.device) – The device that samples should automatically be moved to upon iteration.
prefetch_fn (TPrefetchFn, optional) – A function to run on the data after fetching it. (default: None)

move_to_gpu(batch: composer.core.types.Batch) → composer.core.types.Batch[source]

Move data to the GPU device.

Parameters: batch (Batch) – The data to move the gpu.

normal_iter()[source]: Iterate over data without prefetching.

prefetch_in_cuda_stream_iter()[source]: Iterate over data while prefetching data in a cuda stream.