| |
| from collections.abc import Mapping, Sequence |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data.dataloader import default_collate |
|
|
| from .data_container import DataContainer |
|
|
|
|
| def collate(batch, samples_per_gpu=1): |
| """Puts each data field into a tensor/DataContainer with outer dimension |
| batch size. |
| |
| Extend default_collate to add support for |
| :type:`~mmcv.parallel.DataContainer`. There are 3 cases. |
| |
| 1. cpu_only = True, e.g., meta data |
| 2. cpu_only = False, stack = True, e.g., images tensors |
| 3. cpu_only = False, stack = False, e.g., gt bboxes |
| """ |
|
|
| if not isinstance(batch, Sequence): |
| raise TypeError(f'{batch.dtype} is not supported.') |
|
|
| if isinstance(batch[0], DataContainer): |
| stacked = [] |
| if batch[0].cpu_only: |
| for i in range(0, len(batch), samples_per_gpu): |
| stacked.append( |
| [sample.data for sample in batch[i:i + samples_per_gpu]]) |
| return DataContainer( |
| stacked, batch[0].stack, batch[0].padding_value, cpu_only=True) |
| elif batch[0].stack: |
| for i in range(0, len(batch), samples_per_gpu): |
| assert isinstance(batch[i].data, torch.Tensor) |
|
|
| 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( |
| [sample.data for sample in batch[i:i + samples_per_gpu]]) |
| return DataContainer(stacked, batch[0].stack, batch[0].padding_value) |
| elif isinstance(batch[0], Sequence): |
| transposed = zip(*batch) |
| return [collate(samples, samples_per_gpu) for samples in transposed] |
| elif isinstance(batch[0], Mapping): |
| return { |
| key: collate([d[key] for d in batch], samples_per_gpu) |
| for key in batch[0] |
| } |
| else: |
| return default_collate(batch) |
|
|
