remove main thread clearing for data_queue

apache · Jul 30, 2018 · 7f1be4f · 7f1be4f
1 parent 21f6659
commit 7f1be4f
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 9 deletions.
diff --git a/python/mxnet/gluon/data/dataloader.py b/python/mxnet/gluon/data/dataloader.py
@@ -189,7 +189,7 @@ def __init__(self, num_workers, dataset, batchify_fn, batch_sampler, pin_memory=
         self._batchify_fn = batchify_fn
         self._batch_sampler = batch_sampler
         self._key_queue = Queue()
-        self._data_queue = SimpleQueue()
+        self._data_queue = Queue() if sys.version_info[0] <= 2 else SimpleQueue()
         self._data_buffer = {}
         self._rcvd_idx = 0
         self._sent_idx = 0
@@ -254,11 +254,6 @@ def shutdown(self):
         if not self._shutdown:
             for _ in range(self._num_workers):
                 self._key_queue.put((None, None))
-            try:
-                while not self._data_queue.empty():
-                    self._data_queue.get()
-            except IOError:
-                pass
             self._data_queue.put((None, None))
             self._shutdown = True
 

diff --git a/tests/python/unittest/test_gluon_data.py b/tests/python/unittest/test_gluon_data.py
@@ -202,9 +202,6 @@ def _batchify(data):
 
 @with_seed()
 def test_multi_worker_forked_data_loader():
-    """
-    Test should successfully run its course of multi-process/forked data loader without errors
-    """
     data = _Dummy(False)
     loader = DataLoader(data, batch_size=40, batchify_fn=_batchify, num_workers=2)
     for epoch in range(1):