Polish DataContainer (#79)

* add iminvert

* DC now supports dimension other than 3

* fit flake8

* generalize pad_dim and beautify code

* fix bug; rename "pad_dim" to "pad_dims" for disambiguity

* fix bug

* fix flake8

* fix yapf format

* use deafault yapf

* beautify style

mmcv/parallel/collate.py

@@ -34,21 +34,39 @@ def collate(batch, samples_per_gpu=1):
         elif batch[0].stack:
             for i in range(0, len(batch), samples_per_gpu):
                 assert isinstance(batch[i].data, torch.Tensor)
-                # TODO: handle tensors other than 3d
-                assert batch[i].dim() == 3
-                c, h, w = batch[i].size()
-                for sample in batch[i:i + samples_per_gpu]:
-                    assert c == sample.size(0)
-                    h = max(h, sample.size(1))
-                    w = max(w, sample.size(2))
-                padded_samples = [
-                    F.pad(
-                        sample.data,
-                        (0, w - sample.size(2), 0, h - sample.size(1)),
-                        value=sample.padding_value)
-                    for sample in batch[i:i + samples_per_gpu]
-                ]
-                stacked.append(default_collate(padded_samples))
+
+                if batch[i].pad_dims is not None:
+                    ndim = batch[i].dim()
+                    assert ndim > batch[i].pad_dims
+                    max_shape = [0 for _ in range(batch[i].pad_dims)]
+                    for dim in range(1, batch[i].pad_dims + 1):
+                        max_shape[dim - 1] = batch[i].size(-dim)
+                    for sample in batch[i:i + samples_per_gpu]:
+                        for dim in range(0, ndim - batch[i].pad_dims):
+                            assert batch[i].size(dim) == sample.size(dim)
+                        for dim in range(1, batch[i].pad_dims + 1):
+                            max_shape[dim - 1] = max(max_shape[dim - 1],
+                                                     sample.size(-dim))
+                    padded_samples = []
+                    for sample in batch[i:i + samples_per_gpu]:
+                        pad = [0 for _ in range(batch[i].pad_dims * 2)]
+                        for dim in range(1, batch[i].pad_dims + 1):
+                            pad[2 * dim -
+                                1] = max_shape[dim - 1] - sample.size(-dim)
+                        padded_samples.append(
+                            F.pad(
+                                sample.data, pad, value=sample.padding_value))
+                    stacked.append(default_collate(padded_samples))
+                elif batch[i].pad_dims is None:
+                    stacked.append(
+                        default_collate([
+                            sample.data
+                            for sample in batch[i:i + samples_per_gpu]
+                        ]))
+                else:
+                    raise ValueError(
+                        'pad_dims should be either None or integers (1-3)')
+
         else:
             for i in range(0, len(batch), samples_per_gpu):
                 stacked.append(

mmcv/parallel/data_container.py

@@ -29,13 +29,21 @@ class DataContainer(object):
     - copy to GPU, pad all tensors to the same size and stack them
     - copy to GPU without stacking
     - leave the objects as is and pass it to the model
+    - pad_dims specifies the number of last few dimensions to do padding
     """
 
-    def __init__(self, data, stack=False, padding_value=0, cpu_only=False):
+    def __init__(self,
+                 data,
+                 stack=False,
+                 padding_value=0,
+                 cpu_only=False,
+                 pad_dims=2):
         self._data = data
         self._cpu_only = cpu_only
         self._stack = stack
         self._padding_value = padding_value
+        assert pad_dims in [None, 1, 2, 3]
+        self._pad_dims = pad_dims
 
     def __repr__(self):
         return '{}({})'.format(self.__class__.__name__, repr(self.data))
@@ -63,6 +71,10 @@ class DataContainer(object):
     def padding_value(self):
         return self._padding_value
 
+    @property
+    def pad_dims(self):
+        return self._pad_dims
+
     @assert_tensor_type
     def size(self, *args, **kwargs):
         return self.data.size(*args, **kwargs)