r""""Contains definitions of the methods used by the _BaseDataLoaderIter workers tocollate samples fetched from dataset into Tensor(s).These **needs** to be in global scope since Py2 doesn't support serializingstatic methods.`default_collate` and `default_convert` are exposed to users via 'dataloader.py'."""importtorchimportreimportcollectionsfromtorch._siximportstring_classesnp_str_obj_array_pattern=re.compile(r'[SaUO]')
[docs]defdefault_convert(data):r""" Function that converts each NumPy array element into a :class:`torch.Tensor`. If the input is a `Sequence`, `Collection`, or `Mapping`, it tries to convert each element inside to a :class:`torch.Tensor`. If the input is not an NumPy array, it is left unchanged. This is used as the default function for collation when both `batch_sampler` and `batch_size` are NOT defined in :class:`~torch.utils.data.DataLoader`. The general input type to output type mapping is similar to that of :func:`~torch.utils.data.default_collate`. See the description there for more details. Args: data: a single data point to be converted Examples: >>> # Example with `int` >>> default_convert(0) 0 >>> # Example with NumPy array >>> default_convert(np.array([0, 1])) tensor([0, 1]) >>> # Example with NamedTuple >>> Point = namedtuple('Point', ['x', 'y']) >>> default_convert(Point(0, 0)) Point(x=0, y=0) >>> default_convert(Point(np.array(0), np.array(0))) Point(x=tensor(0), y=tensor(0)) >>> # Example with List >>> default_convert([np.array([0, 1]), np.array([2, 3])]) [tensor([0, 1]), tensor([2, 3])] """elem_type=type(data)ifisinstance(data,torch.Tensor):returndataelifelem_type.__module__=='numpy'andelem_type.__name__!='str_' \
andelem_type.__name__!='string_':# array of string classes and objectifelem_type.__name__=='ndarray' \
andnp_str_obj_array_pattern.search(data.dtype.str)isnotNone:returndatareturntorch.as_tensor(data)elifisinstance(data,collections.abc.Mapping):try:returnelem_type({key:default_convert(data[key])forkeyindata})exceptTypeError:# The mapping type may not support `__init__(iterable)`.return{key:default_convert(data[key])forkeyindata}elifisinstance(data,tuple)andhasattr(data,'_fields'):# namedtuplereturnelem_type(*(default_convert(d)fordindata))elifisinstance(data,tuple):return[default_convert(d)fordindata]# Backwards compatibility.elifisinstance(data,collections.abc.Sequence)andnotisinstance(data,string_classes):try:returnelem_type([default_convert(d)fordindata])exceptTypeError:# The sequence type may not support `__init__(iterable)` (e.g., `range`).return[default_convert(d)fordindata]else:returndata
default_collate_err_msg_format=("default_collate: batch must contain tensors, numpy arrays, numbers, ""dicts or lists; found {}")
[docs]defdefault_collate(batch):r""" Function that takes in a batch of data and puts the elements within the batch into a tensor with an additional outer dimension - batch size. The exact output type can be a :class:`torch.Tensor`, a `Sequence` of :class:`torch.Tensor`, a Collection of :class:`torch.Tensor`, or left unchanged, depending on the input type. This is used as the default function for collation when `batch_size` or `batch_sampler` is defined in :class:`~torch.utils.data.DataLoader`. Here is the general input type (based on the type of the element within the batch) to output type mapping: * :class:`torch.Tensor` -> :class:`torch.Tensor` (with an added outer dimension batch size) * NumPy Arrays -> :class:`torch.Tensor` * `float` -> :class:`torch.Tensor` * `int` -> :class:`torch.Tensor` * `str` -> `str` (unchanged) * `bytes` -> `bytes` (unchanged) * `Mapping[K, V_i]` -> `Mapping[K, default_collate([V_1, V_2, ...])]` * `NamedTuple[V1_i, V2_i, ...]` -> `NamedTuple[default_collate([V1_1, V1_2, ...]), default_collate([V2_1, V2_2, ...]), ...]` * `Sequence[V1_i, V2_i, ...]` -> `Sequence[default_collate([V1_1, V1_2, ...]), default_collate([V2_1, V2_2, ...]), ...]` Args: batch: a single batch to be collated Examples: >>> # Example with a batch of `int`s: >>> default_collate([0, 1, 2, 3]) tensor([0, 1, 2, 3]) >>> # Example with a batch of `str`s: >>> default_collate(['a', 'b', 'c']) ['a', 'b', 'c'] >>> # Example with `Map` inside the batch: >>> default_collate([{'A': 0, 'B': 1}, {'A': 100, 'B': 100}]) {'A': tensor([ 0, 100]), 'B': tensor([ 1, 100])} >>> # Example with `NamedTuple` inside the batch: >>> Point = namedtuple('Point', ['x', 'y']) >>> default_collate([Point(0, 0), Point(1, 1)]) Point(x=tensor([0, 1]), y=tensor([0, 1])) >>> # Example with `Tuple` inside the batch: >>> default_collate([(0, 1), (2, 3)]) [tensor([0, 2]), tensor([1, 3])] >>> # Example with `List` inside the batch: >>> default_collate([[0, 1], [2, 3]]) [tensor([0, 2]), tensor([1, 3])] """elem=batch[0]elem_type=type(elem)ifisinstance(elem,torch.Tensor):out=Noneiftorch.utils.data.get_worker_info()isnotNone:# If we're in a background process, concatenate directly into a# shared memory tensor to avoid an extra copynumel=sum(x.numel()forxinbatch)storage=elem.storage()._new_shared(numel)out=elem.new(storage).resize_(len(batch),*list(elem.size()))returntorch.stack(batch,0,out=out)elifelem_type.__module__=='numpy'andelem_type.__name__!='str_' \
andelem_type.__name__!='string_':ifelem_type.__name__=='ndarray'orelem_type.__name__=='memmap':# array of string classes and objectifnp_str_obj_array_pattern.search(elem.dtype.str)isnotNone:raiseTypeError(default_collate_err_msg_format.format(elem.dtype))returndefault_collate([torch.as_tensor(b)forbinbatch])elifelem.shape==():# scalarsreturntorch.as_tensor(batch)elifisinstance(elem,float):returntorch.tensor(batch,dtype=torch.float64)elifisinstance(elem,int):returntorch.tensor(batch)elifisinstance(elem,string_classes):returnbatchelifisinstance(elem,collections.abc.Mapping):try:returnelem_type({key:default_collate([d[key]fordinbatch])forkeyinelem})exceptTypeError:# The mapping type may not support `__init__(iterable)`.return{key:default_collate([d[key]fordinbatch])forkeyinelem}elifisinstance(elem,tuple)andhasattr(elem,'_fields'):# namedtuplereturnelem_type(*(default_collate(samples)forsamplesinzip(*batch)))elifisinstance(elem,collections.abc.Sequence):# check to make sure that the elements in batch have consistent sizeit=iter(batch)elem_size=len(next(it))ifnotall(len(elem)==elem_sizeforeleminit):raiseRuntimeError('each element in list of batch should be of equal size')transposed=list(zip(*batch))# It may be accessed twice, so we use a list.ifisinstance(elem,tuple):return[default_collate(samples)forsamplesintransposed]# Backwards compatibility.else:try:returnelem_type([default_collate(samples)forsamplesintransposed])exceptTypeError:# The sequence type may not support `__init__(iterable)` (e.g., `range`).return[default_collate(samples)forsamplesintransposed]raiseTypeError(default_collate_err_msg_format.format(elem_type))
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.