samba.sambaloader

SambaLoader

class SambaLoader(dataloader: ~typing.Iterable[~typing.Iterable[~torch.Tensor]], names: ~typing.List[str], maxsize: int = 10, function_hook: ~typing.Callable[[~typing.Iterable[~torch.Tensor]], ~typing.Iterable[~torch.Tensor]] = <function SambaLoader.<lambda>>, return_original_batch: bool = False, ignore_dataloader_exception: bool = False, timeout: float = 900)

New in version 1.19.

A wrapper around PyTorch’s DataLoader class to overlap SambaFlow preprocessing and RDU computations. Similar to the PyTorch DataLoader’s pin_memory feature, SambaLoader inititalizes threads to iterate over the PyTorch DataLoader, perform necessary preprocessing actions like applying transforms registered by the compiler, and send the data to the RDU so the overhead of SambaFlow can be overlapped with the RDU run.

The two parallel threads that run are:

  1. Producer: PyTorch DataLoader -> SambaLoader -> tensor transformations

  2. Consumer: Host-to-device tensor transfers -> chip run -> device-to-host tensor transfers

SambaLoader supports pinned_memory, which is enabled via samba.session.enabled_pinned_memory. Pinned_memory provides buffers that are used to overlap host-to-device tensor transfers and graph run operations. With pinned memory, the two parallel threads would be:

  1. Producer: PyTorch DataLoader -> SambaLoader -> tensor transformations -> host-to-device tensor transfers

  2. Consumer: Chip run -> device-to-host tensor transfers

Example

>>> import torch
>>> from sambaflow.samba.sambaloader import SambaLoader
>>> data = [[torch.tensor(1)], [torch.tensor(2)], [torch.tensor(3)], [torch.tensor(4)]]
>>> names = ["int"]

>>> # looping over SambaLoader
>>> sambaloader = SambaLoader(data, names)
>>> output = []
>>> for item in sambaloader:
...     output.append(item[0].item())
>>> output
[1, 2, 3, 4]

>>> # using SambaLoader as an iterator
>>> sambaloader = SambaLoader(data, names, function_hook=lambda t: [t[0] * -1])
>>> sambaloader_iter = iter(sambaloader)
>>> next(sambaloader_iter)[0].data
tensor(-1)
Parameters:

  • dataloader – the original PyTorch DataLoader

  • names – the names of the input SambaTensor s used when tracing and compiling the PyTorch model

  • maxsize – the maximum number of batches to hold in the buffer for the producer thread.

  • function_hook – a hook for processing the output of the PyTorch DataLoader before the output is handled by SambaFlow. Potential use cases include doing additional transforms on the PyTorch tensors and selectively transferring SambaTensors to device. The SambaLoader still returns all tensors, and the number of tensors returned by the hook must be of the same length as the names argument

  • return_original_batch – if True, returns a tuple of the original batch from the PyTorch DataLoader and the SambaTensor batch from the SambaLoader, i.e. (batch_in_torch, batch_in_samba). Otherwise, returns only the SambaTensor batch from the SambaLoader

  • ignore_dataloader_exception – if True, the SambaLoader proceeds even if there is exception in the PyTorch DataLoader and returns a DataLoaderException object of the PyTorch DataLoader’s exception to the user. Otherwise, SambaLoader raises the exception if the PyTorch DataLoader encounters an error.

  • timeout – how long SambaLoader waits for a batch from the PyTorch DataLoader

exception DataLoaderException

Class for exceptions raised by the original PyTorch DataLoader, (as opposed to SambaLoader). Used when initializing SambaLoader with ignore_dataloader_exception=True.