We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Bug when “numpy.str_” occur in dataloader using default_collate
If, for some reason, my `sample` contains data of type `numpy.str_` instead of `str`, e.g., `type(sample['filepath']) == type(numpy.str_)` then when it passed to `DataLoader`, a `KeyError` would occur. I checked and figure that maybe I found a bug. Consider function `default_collate` below:
```python
def default_collate(batch):
    "Puts each data field into a tensor with outer dimension batch size"
    if torch.is_tensor(batch[0]):
        out = None
        if _use_shared_memory:
            # If we're in a background process, concatenate directly into a
            # shared memory tensor to avoid an extra copy
            numel = sum([x.numel() for x in batch])
            storage = batch[0].storage()._new_shared(numel)
            out = batch[0].new(storage)
        return torch.stack(batch, 0, out=out)
    elif type(batch[0]).__module__ == 'numpy':
        elem = batch[0]
        if type(elem).__name__ == 'ndarray':
            return torch.stack([torch.from_numpy(b) for b in batch], 0)
        if elem.shape == ():  # scalars
            py_type = float if elem.dtype.name.startswith('float') else int
            return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))
    elif isinstance(batch[0], int):
        return torch.LongTensor(batch)
    elif isinstance(batch[0], float):
        return torch.DoubleTensor(batch)
    elif isinstance(batch[0], string_classes):
        return batch
    elif isinstance(batch[0], collections.Mapping):
        return {key: default_collate([d[key] for d in batch]) for key in batch[0]}
    elif isinstance(batch[0], collections.Sequence):
        transposed = zip(*batch)
        return [default_collate(samples) for samples in transposed]

    raise TypeError(("batch must contain tensors, numbers, dicts or lists; found {}"
                     .format(type(batch[0]))))
```
When the data is a `numpy.str_`, it would pass `elif type(batch[0]).__module__ == 'numpy':` and then treated like either a 'ndarray' or a 'float' scale. What I should do is set my `sample['filepath'] = str(sample['filepath'])`, but it would be nice if the function could at least giving more information about the error.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch/utils/data/dataloader.py`
Content:
```
1 import torch
2 import torch.multiprocessing as multiprocessing
3 from .sampler import SequentialSampler, RandomSampler, BatchSampler
4 import collections
5 import sys
6 import traceback
7 import threading
8 from torch._six import string_classes
9 
10 
11 if sys.version_info[0] == 2:
12     import Queue as queue
13 else:
14     import queue
15 
16 
17 _use_shared_memory = False
18 """Whether to use shared memory in default_collate"""
19 
20 
21 class ExceptionWrapper(object):
22     "Wraps an exception plus traceback to communicate across threads"
23 
24     def __init__(self, exc_info):
25         self.exc_type = exc_info[0]
26         self.exc_msg = "".join(traceback.format_exception(*exc_info))
27 
28 
29 def _worker_loop(dataset, index_queue, data_queue, collate_fn):
30     global _use_shared_memory
31     _use_shared_memory = True
32 
33     torch.set_num_threads(1)
34     while True:
35         r = index_queue.get()
36         if r is None:
37             data_queue.put(None)
38             break
39         idx, batch_indices = r
40         try:
41             samples = collate_fn([dataset[i] for i in batch_indices])
42         except Exception:
43             data_queue.put((idx, ExceptionWrapper(sys.exc_info())))
44         else:
45             data_queue.put((idx, samples))
46 
47 
48 def _pin_memory_loop(in_queue, out_queue, done_event):
49     while True:
50         try:
51             r = in_queue.get()
52         except Exception:
53             if done_event.is_set():
54                 return
55             raise
56         if r is None:
57             break
58         if isinstance(r[1], ExceptionWrapper):
59             out_queue.put(r)
60             continue
61         idx, batch = r
62         try:
63             batch = pin_memory_batch(batch)
64         except Exception:
65             out_queue.put((idx, ExceptionWrapper(sys.exc_info())))
66         else:
67             out_queue.put((idx, batch))
68 
69 
70 numpy_type_map = {
71     'float64': torch.DoubleTensor,
72     'float32': torch.FloatTensor,
73     'float16': torch.HalfTensor,
74     'int64': torch.LongTensor,
75     'int32': torch.IntTensor,
76     'int16': torch.ShortTensor,
77     'int8': torch.CharTensor,
78     'uint8': torch.ByteTensor,
79 }
80 
81 
82 def default_collate(batch):
83     "Puts each data field into a tensor with outer dimension batch size"
84     if torch.is_tensor(batch[0]):
85         out = None
86         if _use_shared_memory:
87             # If we're in a background process, concatenate directly into a
88             # shared memory tensor to avoid an extra copy
89             numel = sum([x.numel() for x in batch])
90             storage = batch[0].storage()._new_shared(numel)
91             out = batch[0].new(storage)
92         return torch.stack(batch, 0, out=out)
93     elif type(batch[0]).__module__ == 'numpy':
94         elem = batch[0]
95         if type(elem).__name__ == 'ndarray':
96             return torch.stack([torch.from_numpy(b) for b in batch], 0)
97         if elem.shape == ():  # scalars
98             py_type = float if elem.dtype.name.startswith('float') else int
99             return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))
100     elif isinstance(batch[0], int):
101         return torch.LongTensor(batch)
102     elif isinstance(batch[0], float):
103         return torch.DoubleTensor(batch)
104     elif isinstance(batch[0], string_classes):
105         return batch
106     elif isinstance(batch[0], collections.Mapping):
107         return {key: default_collate([d[key] for d in batch]) for key in batch[0]}
108     elif isinstance(batch[0], collections.Sequence):
109         transposed = zip(*batch)
110         return [default_collate(samples) for samples in transposed]
111 
112     raise TypeError(("batch must contain tensors, numbers, dicts or lists; found {}"
113                      .format(type(batch[0]))))
114 
115 
116 def pin_memory_batch(batch):
117     if torch.is_tensor(batch):
118         return batch.pin_memory()
119     elif isinstance(batch, string_classes):
120         return batch
121     elif isinstance(batch, collections.Mapping):
122         return {k: pin_memory_batch(sample) for k, sample in batch.items()}
123     elif isinstance(batch, collections.Sequence):
124         return [pin_memory_batch(sample) for sample in batch]
125     else:
126         return batch
127 
128 
129 class DataLoaderIter(object):
130     "Iterates once over the DataLoader's dataset, as specified by the sampler"
131 
132     def __init__(self, loader):
133         self.dataset = loader.dataset
134         self.collate_fn = loader.collate_fn
135         self.batch_sampler = loader.batch_sampler
136         self.num_workers = loader.num_workers
137         self.pin_memory = loader.pin_memory
138         self.done_event = threading.Event()
139 
140         self.sample_iter = iter(self.batch_sampler)
141 
142         if self.num_workers > 0:
143             self.index_queue = multiprocessing.SimpleQueue()
144             self.data_queue = multiprocessing.SimpleQueue()
145             self.batches_outstanding = 0
146             self.shutdown = False
147             self.send_idx = 0
148             self.rcvd_idx = 0
149             self.reorder_dict = {}
150 
151             self.workers = [
152                 multiprocessing.Process(
153                     target=_worker_loop,
154                     args=(self.dataset, self.index_queue, self.data_queue, self.collate_fn))
155                 for _ in range(self.num_workers)]
156 
157             for w in self.workers:
158                 w.daemon = True  # ensure that the worker exits on process exit
159                 w.start()
160 
161             if self.pin_memory:
162                 in_data = self.data_queue
163                 self.data_queue = queue.Queue()
164                 self.pin_thread = threading.Thread(
165                     target=_pin_memory_loop,
166                     args=(in_data, self.data_queue, self.done_event))
167                 self.pin_thread.daemon = True
168                 self.pin_thread.start()
169 
170             # prime the prefetch loop
171             for _ in range(2 * self.num_workers):
172                 self._put_indices()
173 
174     def __len__(self):
175         return len(self.batch_sampler)
176 
177     def __next__(self):
178         if self.num_workers == 0:  # same-process loading
179             indices = next(self.sample_iter)  # may raise StopIteration
180             batch = self.collate_fn([self.dataset[i] for i in indices])
181             if self.pin_memory:
182                 batch = pin_memory_batch(batch)
183             return batch
184 
185         # check if the next sample has already been generated
186         if self.rcvd_idx in self.reorder_dict:
187             batch = self.reorder_dict.pop(self.rcvd_idx)
188             return self._process_next_batch(batch)
189 
190         if self.batches_outstanding == 0:
191             self._shutdown_workers()
192             raise StopIteration
193 
194         while True:
195             assert (not self.shutdown and self.batches_outstanding > 0)
196             idx, batch = self.data_queue.get()
197             self.batches_outstanding -= 1
198             if idx != self.rcvd_idx:
199                 # store out-of-order samples
200                 self.reorder_dict[idx] = batch
201                 continue
202             return self._process_next_batch(batch)
203 
204     next = __next__  # Python 2 compatibility
205 
206     def __iter__(self):
207         return self
208 
209     def _put_indices(self):
210         assert self.batches_outstanding < 2 * self.num_workers
211         indices = next(self.sample_iter, None)
212         if indices is None:
213             return
214         self.index_queue.put((self.send_idx, indices))
215         self.batches_outstanding += 1
216         self.send_idx += 1
217 
218     def _process_next_batch(self, batch):
219         self.rcvd_idx += 1
220         self._put_indices()
221         if isinstance(batch, ExceptionWrapper):
222             raise batch.exc_type(batch.exc_msg)
223         return batch
224 
225     def __getstate__(self):
226         # TODO: add limited pickling support for sharing an iterator
227         # across multiple threads for HOGWILD.
228         # Probably the best way to do this is by moving the sample pushing
229         # to a separate thread and then just sharing the data queue
230         # but signalling the end is tricky without a non-blocking API
231         raise NotImplementedError("DataLoaderIterator cannot be pickled")
232 
233     def _shutdown_workers(self):
234         if not self.shutdown:
235             self.shutdown = True
236             self.done_event.set()
237             for _ in self.workers:
238                 self.index_queue.put(None)
239 
240     def __del__(self):
241         if self.num_workers > 0:
242             self._shutdown_workers()
243 
244 
245 class DataLoader(object):
246     """
247     Data loader. Combines a dataset and a sampler, and provides
248     single- or multi-process iterators over the dataset.
249 
250     Arguments:
251         dataset (Dataset): dataset from which to load the data.
252         batch_size (int, optional): how many samples per batch to load
253             (default: 1).
254         shuffle (bool, optional): set to ``True`` to have the data reshuffled
255             at every epoch (default: False).
256         sampler (Sampler, optional): defines the strategy to draw samples from
257             the dataset. If specified, ``shuffle`` must be False.
258         batch_sampler (Sampler, optional): like sampler, but returns a batch of
259             indices at a time. Mutually exclusive with batch_size, shuffle,
260             sampler, and drop_last.
261         num_workers (int, optional): how many subprocesses to use for data
262             loading. 0 means that the data will be loaded in the main process
263             (default: 0)
264         collate_fn (callable, optional): merges a list of samples to form a mini-batch.
265         pin_memory (bool, optional): If ``True``, the data loader will copy tensors
266             into CUDA pinned memory before returning them.
267         drop_last (bool, optional): set to ``True`` to drop the last incomplete batch,
268             if the dataset size is not divisible by the batch size. If False and
269             the size of dataset is not divisible by the batch size, then the last batch
270             will be smaller. (default: False)
271     """
272 
273     def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None,
274                  num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False):
275         self.dataset = dataset
276         self.batch_size = batch_size
277         self.num_workers = num_workers
278         self.collate_fn = collate_fn
279         self.pin_memory = pin_memory
280         self.drop_last = drop_last
281 
282         if batch_sampler is not None:
283             if batch_size > 1 or shuffle or sampler is not None or drop_last:
284                 raise ValueError('batch_sampler is mutually exclusive with '
285                                  'batch_size, shuffle, sampler, and drop_last')
286 
287         if sampler is not None and shuffle:
288             raise ValueError('sampler is mutually exclusive with shuffle')
289 
290         if batch_sampler is None:
291             if sampler is None:
292                 if shuffle:
293                     sampler = RandomSampler(dataset)
294                 else:
295                     sampler = SequentialSampler(dataset)
296             batch_sampler = BatchSampler(sampler, batch_size, drop_last)
297 
298         self.sampler = sampler
299         self.batch_sampler = batch_sampler
300 
301     def __iter__(self):
302         return DataLoaderIter(self)
303 
304     def __len__(self):
305         return len(self.batch_sampler)
306
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
--- a/torch/utils/data/dataloader.py
+++ b/torch/utils/data/dataloader.py
@@ -2,6 +2,7 @@
 import torch.multiprocessing as multiprocessing
 from .sampler import SequentialSampler, RandomSampler, BatchSampler
 import collections
+import re
 import sys
 import traceback
 import threading
@@ -81,6 +82,9 @@
 
 def default_collate(batch):
     "Puts each data field into a tensor with outer dimension batch size"
+
+    error_msg = "batch must contain tensors, numbers, dicts or lists; found {}"
+    elem_type = type(batch[0])
     if torch.is_tensor(batch[0]):
         out = None
         if _use_shared_memory:
@@ -90,9 +94,14 @@
             storage = batch[0].storage()._new_shared(numel)
             out = batch[0].new(storage)
         return torch.stack(batch, 0, out=out)
-    elif type(batch[0]).__module__ == 'numpy':
+    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
+            and elem_type.__name__ != 'string_':
         elem = batch[0]
-        if type(elem).__name__ == 'ndarray':
+        if elem_type.__name__ == 'ndarray':
+            # array of string classes and object
+            if re.search('[SaUO]', elem.dtype.str) is not None:
+                raise TypeError(error_msg.format(elem.dtype))
+
             return torch.stack([torch.from_numpy(b) for b in batch], 0)
         if elem.shape == ():  # scalars
             py_type = float if elem.dtype.name.startswith('float') else int
@@ -109,8 +118,7 @@
         transposed = zip(*batch)
         return [default_collate(samples) for samples in transposed]
 
-    raise TypeError(("batch must contain tensors, numbers, dicts or lists; found {}"
-                     .format(type(batch[0]))))
+    raise TypeError((error_msg.format(type(batch[0]))))
 
 
 def pin_memory_batch(batch):

{"golden_diff": "diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py\n--- a/torch/utils/data/dataloader.py\n+++ b/torch/utils/data/dataloader.py\n@@ -2,6 +2,7 @@\n import torch.multiprocessing as multiprocessing\n from .sampler import SequentialSampler, RandomSampler, BatchSampler\n import collections\n+import re\n import sys\n import traceback\n import threading\n@@ -81,6 +82,9 @@\n \n def default_collate(batch):\n     \"Puts each data field into a tensor with outer dimension batch size\"\n+\n+    error_msg = \"batch must contain tensors, numbers, dicts or lists; found {}\"\n+    elem_type = type(batch[0])\n     if torch.is_tensor(batch[0]):\n         out = None\n         if _use_shared_memory:\n@@ -90,9 +94,14 @@\n             storage = batch[0].storage()._new_shared(numel)\n             out = batch[0].new(storage)\n         return torch.stack(batch, 0, out=out)\n-    elif type(batch[0]).__module__ == 'numpy':\n+    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \\\n+            and elem_type.__name__ != 'string_':\n         elem = batch[0]\n-        if type(elem).__name__ == 'ndarray':\n+        if elem_type.__name__ == 'ndarray':\n+            # array of string classes and object\n+            if re.search('[SaUO]', elem.dtype.str) is not None:\n+                raise TypeError(error_msg.format(elem.dtype))\n+\n             return torch.stack([torch.from_numpy(b) for b in batch], 0)\n         if elem.shape == ():  # scalars\n             py_type = float if elem.dtype.name.startswith('float') else int\n@@ -109,8 +118,7 @@\n         transposed = zip(*batch)\n         return [default_collate(samples) for samples in transposed]\n \n-    raise TypeError((\"batch must contain tensors, numbers, dicts or lists; found {}\"\n-                     .format(type(batch[0]))))\n+    raise TypeError((error_msg.format(type(batch[0]))))\n \n \n def pin_memory_batch(batch):\n", "issue": "Bug when \u201cnumpy.str_\u201d occur in dataloader using default_collate\nIf, for some reason, my `sample` contains data of type `numpy.str_` instead of `str`, e.g., `type(sample['filepath']) == type(numpy.str_)` then when it passed to `DataLoader`, a `KeyError` would occur. I checked and figure that maybe I found a bug. Consider function `default_collate` below:\r\n```python\r\ndef default_collate(batch):\r\n    \"Puts each data field into a tensor with outer dimension batch size\"\r\n    if torch.is_tensor(batch[0]):\r\n        out = None\r\n        if _use_shared_memory:\r\n            # If we're in a background process, concatenate directly into a\r\n            # shared memory tensor to avoid an extra copy\r\n            numel = sum([x.numel() for x in batch])\r\n            storage = batch[0].storage()._new_shared(numel)\r\n            out = batch[0].new(storage)\r\n        return torch.stack(batch, 0, out=out)\r\n    elif type(batch[0]).__module__ == 'numpy':\r\n        elem = batch[0]\r\n        if type(elem).__name__ == 'ndarray':\r\n            return torch.stack([torch.from_numpy(b) for b in batch], 0)\r\n        if elem.shape == ():  # scalars\r\n            py_type = float if elem.dtype.name.startswith('float') else int\r\n            return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))\r\n    elif isinstance(batch[0], int):\r\n        return torch.LongTensor(batch)\r\n    elif isinstance(batch[0], float):\r\n        return torch.DoubleTensor(batch)\r\n    elif isinstance(batch[0], string_classes):\r\n        return batch\r\n    elif isinstance(batch[0], collections.Mapping):\r\n        return {key: default_collate([d[key] for d in batch]) for key in batch[0]}\r\n    elif isinstance(batch[0], collections.Sequence):\r\n        transposed = zip(*batch)\r\n        return [default_collate(samples) for samples in transposed]\r\n\r\n    raise TypeError((\"batch must contain tensors, numbers, dicts or lists; found {}\"\r\n                     .format(type(batch[0]))))\r\n```\r\nWhen the data is a `numpy.str_`, it would pass `elif type(batch[0]).__module__ == 'numpy':` and then treated like either a 'ndarray' or a 'float' scale. What I should do is set my `sample['filepath'] = str(sample['filepath'])`, but it would be nice if the function could at least giving more information about the error.\r\n\n", "before_files": [{"content": "import torch\nimport torch.multiprocessing as multiprocessing\nfrom .sampler import SequentialSampler, RandomSampler, BatchSampler\nimport collections\nimport sys\nimport traceback\nimport threading\nfrom torch._six import string_classes\n\n\nif sys.version_info[0] == 2:\n    import Queue as queue\nelse:\n    import queue\n\n\n_use_shared_memory = False\n\"\"\"Whether to use shared memory in default_collate\"\"\"\n\n\nclass ExceptionWrapper(object):\n    \"Wraps an exception plus traceback to communicate across threads\"\n\n    def __init__(self, exc_info):\n        self.exc_type = exc_info[0]\n        self.exc_msg = \"\".join(traceback.format_exception(*exc_info))\n\n\ndef _worker_loop(dataset, index_queue, data_queue, collate_fn):\n    global _use_shared_memory\n    _use_shared_memory = True\n\n    torch.set_num_threads(1)\n    while True:\n        r = index_queue.get()\n        if r is None:\n            data_queue.put(None)\n            break\n        idx, batch_indices = r\n        try:\n            samples = collate_fn([dataset[i] for i in batch_indices])\n        except Exception:\n            data_queue.put((idx, ExceptionWrapper(sys.exc_info())))\n        else:\n            data_queue.put((idx, samples))\n\n\ndef _pin_memory_loop(in_queue, out_queue, done_event):\n    while True:\n        try:\n            r = in_queue.get()\n        except Exception:\n            if done_event.is_set():\n                return\n            raise\n        if r is None:\n            break\n        if isinstance(r[1], ExceptionWrapper):\n            out_queue.put(r)\n            continue\n        idx, batch = r\n        try:\n            batch = pin_memory_batch(batch)\n        except Exception:\n            out_queue.put((idx, ExceptionWrapper(sys.exc_info())))\n        else:\n            out_queue.put((idx, batch))\n\n\nnumpy_type_map = {\n    'float64': torch.DoubleTensor,\n    'float32': torch.FloatTensor,\n    'float16': torch.HalfTensor,\n    'int64': torch.LongTensor,\n    'int32': torch.IntTensor,\n    'int16': torch.ShortTensor,\n    'int8': torch.CharTensor,\n    'uint8': torch.ByteTensor,\n}\n\n\ndef default_collate(batch):\n    \"Puts each data field into a tensor with outer dimension batch size\"\n    if torch.is_tensor(batch[0]):\n        out = None\n        if _use_shared_memory:\n            # If we're in a background process, concatenate directly into a\n            # shared memory tensor to avoid an extra copy\n            numel = sum([x.numel() for x in batch])\n            storage = batch[0].storage()._new_shared(numel)\n            out = batch[0].new(storage)\n        return torch.stack(batch, 0, out=out)\n    elif type(batch[0]).__module__ == 'numpy':\n        elem = batch[0]\n        if type(elem).__name__ == 'ndarray':\n            return torch.stack([torch.from_numpy(b) for b in batch], 0)\n        if elem.shape == ():  # scalars\n            py_type = float if elem.dtype.name.startswith('float') else int\n            return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))\n    elif isinstance(batch[0], int):\n        return torch.LongTensor(batch)\n    elif isinstance(batch[0], float):\n        return torch.DoubleTensor(batch)\n    elif isinstance(batch[0], string_classes):\n        return batch\n    elif isinstance(batch[0], collections.Mapping):\n        return {key: default_collate([d[key] for d in batch]) for key in batch[0]}\n    elif isinstance(batch[0], collections.Sequence):\n        transposed = zip(*batch)\n        return [default_collate(samples) for samples in transposed]\n\n    raise TypeError((\"batch must contain tensors, numbers, dicts or lists; found {}\"\n                     .format(type(batch[0]))))\n\n\ndef pin_memory_batch(batch):\n    if torch.is_tensor(batch):\n        return batch.pin_memory()\n    elif isinstance(batch, string_classes):\n        return batch\n    elif isinstance(batch, collections.Mapping):\n        return {k: pin_memory_batch(sample) for k, sample in batch.items()}\n    elif isinstance(batch, collections.Sequence):\n        return [pin_memory_batch(sample) for sample in batch]\n    else:\n        return batch\n\n\nclass DataLoaderIter(object):\n    \"Iterates once over the DataLoader's dataset, as specified by the sampler\"\n\n    def __init__(self, loader):\n        self.dataset = loader.dataset\n        self.collate_fn = loader.collate_fn\n        self.batch_sampler = loader.batch_sampler\n        self.num_workers = loader.num_workers\n        self.pin_memory = loader.pin_memory\n        self.done_event = threading.Event()\n\n        self.sample_iter = iter(self.batch_sampler)\n\n        if self.num_workers > 0:\n            self.index_queue = multiprocessing.SimpleQueue()\n            self.data_queue = multiprocessing.SimpleQueue()\n            self.batches_outstanding = 0\n            self.shutdown = False\n            self.send_idx = 0\n            self.rcvd_idx = 0\n            self.reorder_dict = {}\n\n            self.workers = [\n                multiprocessing.Process(\n                    target=_worker_loop,\n                    args=(self.dataset, self.index_queue, self.data_queue, self.collate_fn))\n                for _ in range(self.num_workers)]\n\n            for w in self.workers:\n                w.daemon = True  # ensure that the worker exits on process exit\n                w.start()\n\n            if self.pin_memory:\n                in_data = self.data_queue\n                self.data_queue = queue.Queue()\n                self.pin_thread = threading.Thread(\n                    target=_pin_memory_loop,\n                    args=(in_data, self.data_queue, self.done_event))\n                self.pin_thread.daemon = True\n                self.pin_thread.start()\n\n            # prime the prefetch loop\n            for _ in range(2 * self.num_workers):\n                self._put_indices()\n\n    def __len__(self):\n        return len(self.batch_sampler)\n\n    def __next__(self):\n        if self.num_workers == 0:  # same-process loading\n            indices = next(self.sample_iter)  # may raise StopIteration\n            batch = self.collate_fn([self.dataset[i] for i in indices])\n            if self.pin_memory:\n                batch = pin_memory_batch(batch)\n            return batch\n\n        # check if the next sample has already been generated\n        if self.rcvd_idx in self.reorder_dict:\n            batch = self.reorder_dict.pop(self.rcvd_idx)\n            return self._process_next_batch(batch)\n\n        if self.batches_outstanding == 0:\n            self._shutdown_workers()\n            raise StopIteration\n\n        while True:\n            assert (not self.shutdown and self.batches_outstanding > 0)\n            idx, batch = self.data_queue.get()\n            self.batches_outstanding -= 1\n            if idx != self.rcvd_idx:\n                # store out-of-order samples\n                self.reorder_dict[idx] = batch\n                continue\n            return self._process_next_batch(batch)\n\n    next = __next__  # Python 2 compatibility\n\n    def __iter__(self):\n        return self\n\n    def _put_indices(self):\n        assert self.batches_outstanding < 2 * self.num_workers\n        indices = next(self.sample_iter, None)\n        if indices is None:\n            return\n        self.index_queue.put((self.send_idx, indices))\n        self.batches_outstanding += 1\n        self.send_idx += 1\n\n    def _process_next_batch(self, batch):\n        self.rcvd_idx += 1\n        self._put_indices()\n        if isinstance(batch, ExceptionWrapper):\n            raise batch.exc_type(batch.exc_msg)\n        return batch\n\n    def __getstate__(self):\n        # TODO: add limited pickling support for sharing an iterator\n        # across multiple threads for HOGWILD.\n        # Probably the best way to do this is by moving the sample pushing\n        # to a separate thread and then just sharing the data queue\n        # but signalling the end is tricky without a non-blocking API\n        raise NotImplementedError(\"DataLoaderIterator cannot be pickled\")\n\n    def _shutdown_workers(self):\n        if not self.shutdown:\n            self.shutdown = True\n            self.done_event.set()\n            for _ in self.workers:\n                self.index_queue.put(None)\n\n    def __del__(self):\n        if self.num_workers > 0:\n            self._shutdown_workers()\n\n\nclass DataLoader(object):\n    \"\"\"\n    Data loader. Combines a dataset and a sampler, and provides\n    single- or multi-process iterators over the dataset.\n\n    Arguments:\n        dataset (Dataset): dataset from which to load the data.\n        batch_size (int, optional): how many samples per batch to load\n            (default: 1).\n        shuffle (bool, optional): set to ``True`` to have the data reshuffled\n            at every epoch (default: False).\n        sampler (Sampler, optional): defines the strategy to draw samples from\n            the dataset. If specified, ``shuffle`` must be False.\n        batch_sampler (Sampler, optional): like sampler, but returns a batch of\n            indices at a time. Mutually exclusive with batch_size, shuffle,\n            sampler, and drop_last.\n        num_workers (int, optional): how many subprocesses to use for data\n            loading. 0 means that the data will be loaded in the main process\n            (default: 0)\n        collate_fn (callable, optional): merges a list of samples to form a mini-batch.\n        pin_memory (bool, optional): If ``True``, the data loader will copy tensors\n            into CUDA pinned memory before returning them.\n        drop_last (bool, optional): set to ``True`` to drop the last incomplete batch,\n            if the dataset size is not divisible by the batch size. If False and\n            the size of dataset is not divisible by the batch size, then the last batch\n            will be smaller. (default: False)\n    \"\"\"\n\n    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None,\n                 num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False):\n        self.dataset = dataset\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n        self.collate_fn = collate_fn\n        self.pin_memory = pin_memory\n        self.drop_last = drop_last\n\n        if batch_sampler is not None:\n            if batch_size > 1 or shuffle or sampler is not None or drop_last:\n                raise ValueError('batch_sampler is mutually exclusive with '\n                                 'batch_size, shuffle, sampler, and drop_last')\n\n        if sampler is not None and shuffle:\n            raise ValueError('sampler is mutually exclusive with shuffle')\n\n        if batch_sampler is None:\n            if sampler is None:\n                if shuffle:\n                    sampler = RandomSampler(dataset)\n                else:\n                    sampler = SequentialSampler(dataset)\n            batch_sampler = BatchSampler(sampler, batch_size, drop_last)\n\n        self.sampler = sampler\n        self.batch_sampler = batch_sampler\n\n    def __iter__(self):\n        return DataLoaderIter(self)\n\n    def __len__(self):\n        return len(self.batch_sampler)\n", "path": "torch/utils/data/dataloader.py"}], "after_files": [{"content": "import torch\nimport torch.multiprocessing as multiprocessing\nfrom .sampler import SequentialSampler, RandomSampler, BatchSampler\nimport collections\nimport re\nimport sys\nimport traceback\nimport threading\nfrom torch._six import string_classes\n\n\nif sys.version_info[0] == 2:\n    import Queue as queue\nelse:\n    import queue\n\n\n_use_shared_memory = False\n\"\"\"Whether to use shared memory in default_collate\"\"\"\n\n\nclass ExceptionWrapper(object):\n    \"Wraps an exception plus traceback to communicate across threads\"\n\n    def __init__(self, exc_info):\n        self.exc_type = exc_info[0]\n        self.exc_msg = \"\".join(traceback.format_exception(*exc_info))\n\n\ndef _worker_loop(dataset, index_queue, data_queue, collate_fn):\n    global _use_shared_memory\n    _use_shared_memory = True\n\n    torch.set_num_threads(1)\n    while True:\n        r = index_queue.get()\n        if r is None:\n            data_queue.put(None)\n            break\n        idx, batch_indices = r\n        try:\n            samples = collate_fn([dataset[i] for i in batch_indices])\n        except Exception:\n            data_queue.put((idx, ExceptionWrapper(sys.exc_info())))\n        else:\n            data_queue.put((idx, samples))\n\n\ndef _pin_memory_loop(in_queue, out_queue, done_event):\n    while True:\n        try:\n            r = in_queue.get()\n        except Exception:\n            if done_event.is_set():\n                return\n            raise\n        if r is None:\n            break\n        if isinstance(r[1], ExceptionWrapper):\n            out_queue.put(r)\n            continue\n        idx, batch = r\n        try:\n            batch = pin_memory_batch(batch)\n        except Exception:\n            out_queue.put((idx, ExceptionWrapper(sys.exc_info())))\n        else:\n            out_queue.put((idx, batch))\n\n\nnumpy_type_map = {\n    'float64': torch.DoubleTensor,\n    'float32': torch.FloatTensor,\n    'float16': torch.HalfTensor,\n    'int64': torch.LongTensor,\n    'int32': torch.IntTensor,\n    'int16': torch.ShortTensor,\n    'int8': torch.CharTensor,\n    'uint8': torch.ByteTensor,\n}\n\n\ndef default_collate(batch):\n    \"Puts each data field into a tensor with outer dimension batch size\"\n\n    error_msg = \"batch must contain tensors, numbers, dicts or lists; found {}\"\n    elem_type = type(batch[0])\n    if torch.is_tensor(batch[0]):\n        out = None\n        if _use_shared_memory:\n            # If we're in a background process, concatenate directly into a\n            # shared memory tensor to avoid an extra copy\n            numel = sum([x.numel() for x in batch])\n            storage = batch[0].storage()._new_shared(numel)\n            out = batch[0].new(storage)\n        return torch.stack(batch, 0, out=out)\n    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \\\n            and elem_type.__name__ != 'string_':\n        elem = batch[0]\n        if elem_type.__name__ == 'ndarray':\n            # array of string classes and object\n            if re.search('[SaUO]', elem.dtype.str) is not None:\n                raise TypeError(error_msg.format(elem.dtype))\n\n            return torch.stack([torch.from_numpy(b) for b in batch], 0)\n        if elem.shape == ():  # scalars\n            py_type = float if elem.dtype.name.startswith('float') else int\n            return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))\n    elif isinstance(batch[0], int):\n        return torch.LongTensor(batch)\n    elif isinstance(batch[0], float):\n        return torch.DoubleTensor(batch)\n    elif isinstance(batch[0], string_classes):\n        return batch\n    elif isinstance(batch[0], collections.Mapping):\n        return {key: default_collate([d[key] for d in batch]) for key in batch[0]}\n    elif isinstance(batch[0], collections.Sequence):\n        transposed = zip(*batch)\n        return [default_collate(samples) for samples in transposed]\n\n    raise TypeError((error_msg.format(type(batch[0]))))\n\n\ndef pin_memory_batch(batch):\n    if torch.is_tensor(batch):\n        return batch.pin_memory()\n    elif isinstance(batch, string_classes):\n        return batch\n    elif isinstance(batch, collections.Mapping):\n        return {k: pin_memory_batch(sample) for k, sample in batch.items()}\n    elif isinstance(batch, collections.Sequence):\n        return [pin_memory_batch(sample) for sample in batch]\n    else:\n        return batch\n\n\nclass DataLoaderIter(object):\n    \"Iterates once over the DataLoader's dataset, as specified by the sampler\"\n\n    def __init__(self, loader):\n        self.dataset = loader.dataset\n        self.collate_fn = loader.collate_fn\n        self.batch_sampler = loader.batch_sampler\n        self.num_workers = loader.num_workers\n        self.pin_memory = loader.pin_memory\n        self.done_event = threading.Event()\n\n        self.sample_iter = iter(self.batch_sampler)\n\n        if self.num_workers > 0:\n            self.index_queue = multiprocessing.SimpleQueue()\n            self.data_queue = multiprocessing.SimpleQueue()\n            self.batches_outstanding = 0\n            self.shutdown = False\n            self.send_idx = 0\n            self.rcvd_idx = 0\n            self.reorder_dict = {}\n\n            self.workers = [\n                multiprocessing.Process(\n                    target=_worker_loop,\n                    args=(self.dataset, self.index_queue, self.data_queue, self.collate_fn))\n                for _ in range(self.num_workers)]\n\n            for w in self.workers:\n                w.daemon = True  # ensure that the worker exits on process exit\n                w.start()\n\n            if self.pin_memory:\n                in_data = self.data_queue\n                self.data_queue = queue.Queue()\n                self.pin_thread = threading.Thread(\n                    target=_pin_memory_loop,\n                    args=(in_data, self.data_queue, self.done_event))\n                self.pin_thread.daemon = True\n                self.pin_thread.start()\n\n            # prime the prefetch loop\n            for _ in range(2 * self.num_workers):\n                self._put_indices()\n\n    def __len__(self):\n        return len(self.batch_sampler)\n\n    def __next__(self):\n        if self.num_workers == 0:  # same-process loading\n            indices = next(self.sample_iter)  # may raise StopIteration\n            batch = self.collate_fn([self.dataset[i] for i in indices])\n            if self.pin_memory:\n                batch = pin_memory_batch(batch)\n            return batch\n\n        # check if the next sample has already been generated\n        if self.rcvd_idx in self.reorder_dict:\n            batch = self.reorder_dict.pop(self.rcvd_idx)\n            return self._process_next_batch(batch)\n\n        if self.batches_outstanding == 0:\n            self._shutdown_workers()\n            raise StopIteration\n\n        while True:\n            assert (not self.shutdown and self.batches_outstanding > 0)\n            idx, batch = self.data_queue.get()\n            self.batches_outstanding -= 1\n            if idx != self.rcvd_idx:\n                # store out-of-order samples\n                self.reorder_dict[idx] = batch\n                continue\n            return self._process_next_batch(batch)\n\n    next = __next__  # Python 2 compatibility\n\n    def __iter__(self):\n        return self\n\n    def _put_indices(self):\n        assert self.batches_outstanding < 2 * self.num_workers\n        indices = next(self.sample_iter, None)\n        if indices is None:\n            return\n        self.index_queue.put((self.send_idx, indices))\n        self.batches_outstanding += 1\n        self.send_idx += 1\n\n    def _process_next_batch(self, batch):\n        self.rcvd_idx += 1\n        self._put_indices()\n        if isinstance(batch, ExceptionWrapper):\n            raise batch.exc_type(batch.exc_msg)\n        return batch\n\n    def __getstate__(self):\n        # TODO: add limited pickling support for sharing an iterator\n        # across multiple threads for HOGWILD.\n        # Probably the best way to do this is by moving the sample pushing\n        # to a separate thread and then just sharing the data queue\n        # but signalling the end is tricky without a non-blocking API\n        raise NotImplementedError(\"DataLoaderIterator cannot be pickled\")\n\n    def _shutdown_workers(self):\n        if not self.shutdown:\n            self.shutdown = True\n            self.done_event.set()\n            for _ in self.workers:\n                self.index_queue.put(None)\n\n    def __del__(self):\n        if self.num_workers > 0:\n            self._shutdown_workers()\n\n\nclass DataLoader(object):\n    \"\"\"\n    Data loader. Combines a dataset and a sampler, and provides\n    single- or multi-process iterators over the dataset.\n\n    Arguments:\n        dataset (Dataset): dataset from which to load the data.\n        batch_size (int, optional): how many samples per batch to load\n            (default: 1).\n        shuffle (bool, optional): set to ``True`` to have the data reshuffled\n            at every epoch (default: False).\n        sampler (Sampler, optional): defines the strategy to draw samples from\n            the dataset. If specified, ``shuffle`` must be False.\n        batch_sampler (Sampler, optional): like sampler, but returns a batch of\n            indices at a time. Mutually exclusive with batch_size, shuffle,\n            sampler, and drop_last.\n        num_workers (int, optional): how many subprocesses to use for data\n            loading. 0 means that the data will be loaded in the main process\n            (default: 0)\n        collate_fn (callable, optional): merges a list of samples to form a mini-batch.\n        pin_memory (bool, optional): If ``True``, the data loader will copy tensors\n            into CUDA pinned memory before returning them.\n        drop_last (bool, optional): set to ``True`` to drop the last incomplete batch,\n            if the dataset size is not divisible by the batch size. If False and\n            the size of dataset is not divisible by the batch size, then the last batch\n            will be smaller. (default: False)\n    \"\"\"\n\n    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None,\n                 num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False):\n        self.dataset = dataset\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n        self.collate_fn = collate_fn\n        self.pin_memory = pin_memory\n        self.drop_last = drop_last\n\n        if batch_sampler is not None:\n            if batch_size > 1 or shuffle or sampler is not None or drop_last:\n                raise ValueError('batch_sampler is mutually exclusive with '\n                                 'batch_size, shuffle, sampler, and drop_last')\n\n        if sampler is not None and shuffle:\n            raise ValueError('sampler is mutually exclusive with shuffle')\n\n        if batch_sampler is None:\n            if sampler is None:\n                if shuffle:\n                    sampler = RandomSampler(dataset)\n                else:\n                    sampler = SequentialSampler(dataset)\n            batch_sampler = BatchSampler(sampler, batch_size, drop_last)\n\n        self.sampler = sampler\n        self.batch_sampler = batch_sampler\n\n    def __iter__(self):\n        return DataLoaderIter(self)\n\n    def __len__(self):\n        return len(self.batch_sampler)\n", "path": "torch/utils/data/dataloader.py"}]}