MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / Collator

Class Collator

imperative/python/megengine/data/collator.py:28–64  ·  view source on GitHub ↗

r"""Used for merging a list of samples to form a mini-batch of Tensor(s). Used when using batched loading from a dataset. Modified from https://github.com/pytorch/pytorch/blob/master/torch/utils/data/_utils/collate.py

Source from the content-addressed store, hash-verified

26
27
28class Collator:
29 r"""Used for merging a list of samples to form a mini-batch of Tensor(s). Used when using batched loading from a dataset.
30 Modified from https://github.com/pytorch/pytorch/blob/master/torch/utils/data/_utils/collate.py
31 """
32
33 def apply(self, inputs):
34 elem = inputs[0]
35 elem_type = type(elem)
36 if (
37 elem_type.__module__ == "numpy"
38 and elem_type.__name__ != "str_"
39 and elem_type.__name__ != "string_"
40 ):
41 elem = inputs[0]
42 if elem_type.__name__ == "ndarray":
43 # array of string classes and object
44 if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
45 raise TypeError(default_collate_err_msg_format.format(elem.dtype))
46
47 return np.ascontiguousarray(np.stack(inputs))
48 elif elem.shape == (): # scalars
49 return np.array(inputs)
50 elif isinstance(elem, float):
51 return np.array(inputs, dtype=np.float64)
52 elif isinstance(elem, int):
53 return np.array(inputs)
54 elif isinstance(elem, (str, bytes)):
55 return inputs
56 elif isinstance(elem, collections.abc.Mapping):
57 return {key: self.apply([d[key] for d in inputs]) for key in elem}
58 elif isinstance(elem, tuple) and hasattr(elem, "_fields"): # namedtuple
59 return elem_type(*(self.apply(samples) for samples in zip(*inputs)))
60 elif isinstance(elem, collections.abc.Sequence):
61 transposed = zip(*inputs)
62 return [self.apply(samples) for samples in transposed]
63
64 raise TypeError(default_collate_err_msg_format.format(elem_type))

Callers 1

__init__Method · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected