Take a list of samples (as dictionary) and create a batch, preserving the keys. If `tensors` is True, `ndarray` objects are combined into tensor batches. :param dict samples: list of samples :param bool tensors: whether to turn lists of ndarrays into a single ndarray :returns: s
(samples, combine_tensors=True, combine_scalars=True)
| 17 | |
| 18 | |
| 19 | def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True): |
| 20 | """Take a list of samples (as dictionary) and create a batch, preserving the keys. |
| 21 | If `tensors` is True, `ndarray` objects are combined into |
| 22 | tensor batches. |
| 23 | :param dict samples: list of samples |
| 24 | :param bool tensors: whether to turn lists of ndarrays into a single ndarray |
| 25 | :returns: single sample consisting of a batch |
| 26 | :rtype: dict |
| 27 | """ |
| 28 | keys = set.intersection(*[set(sample.keys()) for sample in samples]) |
| 29 | batched = {key: [] for key in keys} |
| 30 | |
| 31 | for s in samples: |
| 32 | [batched[key].append(s[key]) for key in batched] |
| 33 | |
| 34 | result = {} |
| 35 | for key in batched: |
| 36 | if isinstance(batched[key][0], (int, float)): |
| 37 | if combine_scalars: |
| 38 | result[key] = np.array(list(batched[key])) |
| 39 | elif isinstance(batched[key][0], torch.Tensor): |
| 40 | if combine_tensors: |
| 41 | result[key] = torch.stack(list(batched[key])) |
| 42 | elif isinstance(batched[key][0], np.ndarray): |
| 43 | if combine_tensors: |
| 44 | result[key] = np.array(list(batched[key])) |
| 45 | else: |
| 46 | result[key] = list(batched[key]) |
| 47 | return result |
| 48 | |
| 49 | |
| 50 | def identity(x): |
nothing calls this directly
no outgoing calls
no test coverage detected