| 621 | |
| 622 | |
| 623 | class _PrefetchingIter(object): |
| 624 | def __init__(self, dataloader, dataloader_it, num_threads=None): |
| 625 | self.queue = Queue(1) |
| 626 | self.dataloader_it = dataloader_it |
| 627 | self.dataloader = dataloader |
| 628 | self.num_threads = num_threads |
| 629 | |
| 630 | self.use_thread = dataloader.use_prefetch_thread |
| 631 | self.use_alternate_streams = dataloader.use_alternate_streams |
| 632 | self.device = self.dataloader.device |
| 633 | if self.use_alternate_streams and self.device.type == "cuda": |
| 634 | self.stream = torch.cuda.Stream(device=self.device) |
| 635 | else: |
| 636 | self.stream = None |
| 637 | self._shutting_down = False |
| 638 | if self.use_thread: |
| 639 | self._done_event = threading.Event() |
| 640 | thread = threading.Thread( |
| 641 | target=_prefetcher_entry, |
| 642 | args=( |
| 643 | dataloader_it, |
| 644 | dataloader, |
| 645 | self.queue, |
| 646 | num_threads, |
| 647 | self.stream, |
| 648 | self._done_event, |
| 649 | ), |
| 650 | daemon=True, |
| 651 | ) |
| 652 | thread.start() |
| 653 | self.thread = thread |
| 654 | |
| 655 | def __iter__(self): |
| 656 | return self |
| 657 | |
| 658 | def _shutdown(self): |
| 659 | # Sometimes when Python is exiting complicated operations like |
| 660 | # self.queue.get_nowait() will hang. So we set it to no-op and let Python handle |
| 661 | # the rest since the thread is daemonic. |
| 662 | # PyTorch takes the same solution. |
| 663 | if PYTHON_EXIT_STATUS is True or PYTHON_EXIT_STATUS is None: |
| 664 | return |
| 665 | if not self._shutting_down: |
| 666 | try: |
| 667 | self._shutting_down = True |
| 668 | self._done_event.set() |
| 669 | |
| 670 | try: |
| 671 | self.queue.get_nowait() # In case the thread is blocking on put(). |
| 672 | except: # pylint: disable=bare-except |
| 673 | pass |
| 674 | |
| 675 | self.thread.join() |
| 676 | except: # pylint: disable=bare-except |
| 677 | pass |
| 678 | |
| 679 | def __del__(self): |
| 680 | if self.use_thread: |