Method __init__

hanlp/common/dataset.py:553–581

A dataloader wrapper which speeds up bachifying using multi-processing. It works best for dataloaders of which the bachify takes very long time. But it introduces extra GPU memory consumption since prefetched batches are stored in a ``Queue`` on GPU. .. Caution::

(self, dataloader: torch.utils.data.DataLoader, prefetch: int = 10, batchify: Callable = None)

class PrefetchDataLoader(DataLoader):
    def __init__(self, dataloader: torch.utils.data.DataLoader, prefetch: int = 10, batchify: Callable = None) -> None:
        """ A dataloader wrapper which speeds up bachifying using multi-processing. It works best for dataloaders
        of which the bachify takes very long time. But it introduces extra GPU memory consumption since prefetched
        batches are stored in a ``Queue`` on GPU.

        .. Caution::

            PrefetchDataLoader only works in spawn mode with the following initialization code:

            Examples::

                if __name__ == '__main__':
                    import torch

                    torch.multiprocessing.set_start_method('spawn')

            And these 2 lines **MUST** be put into ``if __name__ == '__main__':`` block.

        Args:
            dataloader: A :class:`~torch.utils.data.DatasetLoader` to be prefetched.
            prefetch: Number of batches to prefetch.
            batchify: A bachify function called on each batch of samples. In which case, the inner dataloader shall
                    return samples without really bachify them.
        """
        super().__init__(dataset=dataloader)
        self._batchify = batchify
        self.prefetch = None if isdebugging() else prefetch
        if self.prefetch:
            self._fire_process(dataloader, prefetch)

    def _fire_process(self, dataloader, prefetch):
        self.queue = mp.Queue(prefetch)