Add pin_device_id option to Gluon DataLoader (apache#14136)

* add pin_device_id option to DataLoader

* add unit test to check output context

* trigger CI

python/mxnet/gluon/data/dataloader.py

@@ -169,14 +169,15 @@ def worker_loop_v1(dataset, key_queue, data_queue, batchify_fn):
         batch = batchify_fn([dataset[i] for i in samples])
         data_queue.put((idx, batch))
 
-def fetcher_loop_v1(data_queue, data_buffer, pin_memory=False, data_buffer_lock=None):
+def fetcher_loop_v1(data_queue, data_buffer, pin_memory=False,
+                    pin_device_id=0, data_buffer_lock=None):
     """Fetcher loop for fetching data from queue and put in reorder dict."""
     while True:
         idx, batch = data_queue.get()
         if idx is None:
             break
         if pin_memory:
-            batch = _as_in_context(batch, context.cpu_pinned())
+            batch = _as_in_context(batch, context.cpu_pinned(pin_device_id))
         else:
             batch = _as_in_context(batch, context.cpu())
         if data_buffer_lock is not None:
@@ -188,8 +189,8 @@ def fetcher_loop_v1(data_queue, data_buffer, pin_memory=False, data_buffer_lock=
 
 class _MultiWorkerIterV1(object):
     """Internal multi-worker iterator for DataLoader."""
-    def __init__(self, num_workers, dataset, batchify_fn, batch_sampler, pin_memory=False,
-                 worker_fn=worker_loop_v1):
+    def __init__(self, num_workers, dataset, batchify_fn, batch_sampler,
+                 pin_memory=False, pin_device_id=0, worker_fn=worker_loop_v1):
         assert num_workers > 0, "_MultiWorkerIter is not for {} workers".format(num_workers)
         self._num_workers = num_workers
         self._dataset = dataset
@@ -218,7 +219,8 @@ def __init__(self, num_workers, dataset, batchify_fn, batch_sampler, pin_memory=
 
         self._fetcher = threading.Thread(
             target=fetcher_loop_v1,
-            args=(self._data_queue, self._data_buffer, pin_memory, self._data_buffer_lock))
+            args=(self._data_queue, self._data_buffer, pin_memory,
+                  pin_device_id, self._data_buffer_lock))
         self._fetcher.daemon = True
         self._fetcher.start()
 
@@ -323,12 +325,15 @@ def default_batchify_fn(data):
         If ``True``, the dataloader will copy NDArrays into pinned memory
         before returning them. Copying from CPU pinned memory to GPU is faster
         than from normal CPU memory.
+    pin_device_id : int, default 0
+        The device id to use for allocating pinned memory if pin_memory is ``True``
     """
     def __init__(self, dataset, batch_size=None, shuffle=False, sampler=None,
                  last_batch=None, batch_sampler=None, batchify_fn=None,
-                 num_workers=0, pin_memory=False):
+                 num_workers=0, pin_memory=False, pin_device_id=0):
         self._dataset = dataset
         self._pin_memory = pin_memory
+        self._pin_device_id = pin_device_id
 
         if batch_sampler is None:
             if batch_size is None:
@@ -365,13 +370,14 @@ def same_process_iter():
                 for batch in self._batch_sampler:
                     ret = self._batchify_fn([self._dataset[idx] for idx in batch])
                     if self._pin_memory:
-                        ret = _as_in_context(ret, context.cpu_pinned())
+                        ret = _as_in_context(ret, context.cpu_pinned(self._pin_device_id))
                     yield ret
             return same_process_iter()
 
         # multi-worker
         return _MultiWorkerIterV1(self._num_workers, self._dataset,
-                                  self._batchify_fn, self._batch_sampler, self._pin_memory)
+                                  self._batchify_fn, self._batch_sampler,
+                                  self._pin_memory, self._pin_device_id)
 
     def __len__(self):
         return len(self._batch_sampler)
@@ -403,7 +409,7 @@ def _thread_worker_fn(samples, batchify_fn, dataset):
 class _MultiWorkerIter(object):
     """Internal multi-worker iterator for DataLoader."""
     def __init__(self, worker_pool, batchify_fn, batch_sampler, pin_memory=False,
-                 worker_fn=_worker_fn, prefetch=0, dataset=None):
+                 pin_device_id=0, worker_fn=_worker_fn, prefetch=0, dataset=None):
         self._worker_pool = worker_pool
         self._batchify_fn = batchify_fn
         self._batch_sampler = batch_sampler
@@ -413,6 +419,7 @@ def __init__(self, worker_pool, batchify_fn, batch_sampler, pin_memory=False,
         self._iter = iter(self._batch_sampler)
         self._worker_fn = worker_fn
         self._pin_memory = pin_memory
+        self._pin_device_id = pin_device_id
         self._dataset = dataset
         # pre-fetch
         for _ in range(prefetch):
@@ -442,7 +449,7 @@ def __next__(self):
         ret = self._data_buffer.pop(self._rcvd_idx)
         batch = pickle.loads(ret.get()) if self._dataset is None else ret.get()
         if self._pin_memory:
-            batch = _as_in_context(batch, context.cpu_pinned())
+            batch = _as_in_context(batch, context.cpu_pinned(self._pin_device_id))
         batch = batch[0] if len(batch) == 1 else batch
         self._rcvd_idx += 1
         return batch
@@ -498,6 +505,8 @@ def default_batchify_fn(data):
         If ``True``, the dataloader will copy NDArrays into pinned memory
         before returning them. Copying from CPU pinned memory to GPU is faster
         than from normal CPU memory.
+    pin_device_id : int, default 0
+        The device id to use for allocating pinned memory if pin_memory is ``True``
     prefetch : int, default is `num_workers * 2`
         The number of prefetching batches only works if `num_workers` > 0.
         If `prefetch` > 0, it allow worker process to prefetch certain batches before
@@ -514,9 +523,11 @@ def default_batchify_fn(data):
     """
     def __init__(self, dataset, batch_size=None, shuffle=False, sampler=None,
                  last_batch=None, batch_sampler=None, batchify_fn=None,
-                 num_workers=0, pin_memory=False, prefetch=None, thread_pool=False):
+                 num_workers=0, pin_memory=False, pin_device_id=0,
+                 prefetch=None, thread_pool=False):
         self._dataset = dataset
         self._pin_memory = pin_memory
+        self._pin_device_id = pin_device_id
         self._thread_pool = thread_pool
 
         if batch_sampler is None:
@@ -562,13 +573,13 @@ def same_process_iter():
                 for batch in self._batch_sampler:
                     ret = self._batchify_fn([self._dataset[idx] for idx in batch])
                     if self._pin_memory:
-                        ret = _as_in_context(ret, context.cpu_pinned())
+                        ret = _as_in_context(ret, context.cpu_pinned(self._pin_device_id))
                     yield ret
             return same_process_iter()
 
         # multi-worker
         return _MultiWorkerIter(self._worker_pool, self._batchify_fn, self._batch_sampler,
-                                pin_memory=self._pin_memory,
+                                pin_memory=self._pin_memory, pin_device_id=self._pin_device_id,
                                 worker_fn=_thread_worker_fn if self._thread_pool else _worker_fn,
                                 prefetch=self._prefetch,
                                 dataset=self._dataset if self._thread_pool else None)

tests/python/unittest/test_gluon_data.py

@@ -256,6 +256,30 @@ def test_multi_worker_dataloader_release_pool():
         del the_iter
         del D
 
+
+def test_dataloader_context():
+    X = np.random.uniform(size=(10, 20))
+    dataset = gluon.data.ArrayDataset(X)
+    default_dev_id = 0
+    custom_dev_id = 1
+
+    # use non-pinned memory
+    loader1 = gluon.data.DataLoader(dataset, 8)
+    for _, x in enumerate(loader1):
+        assert x.context == context.cpu(default_dev_id)
+
+    # use pinned memory with default device id
+    loader2 = gluon.data.DataLoader(dataset, 8, pin_memory=True)
+    for _, x in enumerate(loader2):
+        assert x.context == context.cpu_pinned(default_dev_id)
+
+    # use pinned memory with custom device id
+    loader3 = gluon.data.DataLoader(dataset, 8, pin_memory=True,
+                                    pin_device_id=custom_dev_id)
+    for _, x in enumerate(loader3):
+        assert x.context == context.cpu_pinned(custom_dev_id)
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()