Repository: pytorch/pytorch

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.



Discussion:
we should give a better error message here. thanks for reporting.
Thanks

Resolving patch:
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 _pin_memory_loop(in_queue, out_queue, done_event):
 
 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 @@ def default_collate(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':
+    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 @@ def default_collate(batch):
         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):


Test patch:
diff --git a/test/test_dataloader.py b/test/test_dataloader.py
--- a/test/test_dataloader.py
+++ b/test/test_dataloader.py
@@ -4,6 +4,7 @@
 import traceback
 import unittest
 from torch.utils.data import Dataset, TensorDataset, DataLoader, ConcatDataset
+from torch.utils.data.dataloader import default_collate
 from common import TestCase, run_tests, TEST_NUMPY
 from common_nn import TEST_CUDA
 
@@ -276,6 +277,23 @@ def __len__(self):
             batch = next(iter(loader))
             self.assertIsInstance(batch, tt)
 
+    @unittest.skipIf(not TEST_NUMPY, "numpy unavailable")
+    def test_default_colate_bad_numpy_types(self):
+        import numpy as np
+
+        # Should be a no-op
+        arr = np.array(['a', 'b', 'c'])
+        default_collate(arr)
+
+        arr = np.array([[['a', 'b', 'c']]])
+        self.assertRaises(TypeError, lambda: default_collate(arr))
+
+        arr = np.array([object(), object(), object()])
+        self.assertRaises(TypeError, lambda: default_collate(arr))
+
+        arr = np.array([[[object(), object(), object()]]])
+        self.assertRaises(TypeError, lambda: default_collate(arr))
+
 
 class StringDataset(Dataset):
     def __init__(self):