chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,36 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .batch_sampler import ( # noqa: F401
|
||||
BatchSampler,
|
||||
DistributedBatchSampler,
|
||||
)
|
||||
from .dataset import ( # noqa: F401
|
||||
ChainDataset,
|
||||
ComposeDataset,
|
||||
ConcatDataset,
|
||||
Dataset,
|
||||
IterableDataset,
|
||||
Subset,
|
||||
TensorDataset,
|
||||
random_split,
|
||||
)
|
||||
from .sampler import ( # noqa: F401
|
||||
RandomSampler,
|
||||
Sampler,
|
||||
SequenceSampler,
|
||||
SubsetRandomSampler,
|
||||
WeightedRandomSampler,
|
||||
)
|
||||
from .worker import get_worker_info # noqa: F401
|
||||
@@ -0,0 +1,430 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Iterable, Iterator, Sequence, Sized
|
||||
from typing import overload
|
||||
|
||||
import numpy as np
|
||||
|
||||
from paddle.utils.decorator_utils import batch_sampler_decorator
|
||||
|
||||
from .dataset import IterableDataset
|
||||
from .sampler import RandomSampler, Sampler, SequenceSampler
|
||||
|
||||
|
||||
class BatchSampler(Sampler[Sequence[int]]):
|
||||
"""
|
||||
A base implement of batch sampler used by `paddle.io.DataLoader`
|
||||
which yield mini-batch indices(a list/tuple with length as
|
||||
mini-batch size and holds sample indices) iterably.
|
||||
|
||||
Batch sampler used by :code:`paddle.io.DataLoader` should be a subclass
|
||||
of :code:`paddle.io.BatchSampler`, BatchSampler subclasses should
|
||||
implement following methods:
|
||||
|
||||
:code:`__iter__`: return mini-batch indices iterably.
|
||||
|
||||
:code:`__len__`: get mini-batch number in an epoch.
|
||||
|
||||
|
||||
Args:
|
||||
dataset(Dataset, optional): this should be an instance of a subclass of :ref:`api_paddle_io_Dataset` or
|
||||
:ref:`api_paddle_io_IterableDataset` or other python object which implemented
|
||||
:code:`__len__` for BatchSampler to get indices as the
|
||||
range of :attr:`dataset` length. Default None, disabled.
|
||||
sampler (Sampler, Iterable, optional): this should be a :ref:`api_paddle_io_Sample` or Iterable
|
||||
instance which implemented :code:`__iter__` to generate
|
||||
sample indices. :attr:`sampler` and :attr:`dataset`
|
||||
can not be set in the same time. If :attr:`sampler`
|
||||
is set, :attr:`dataset` should not be set. Default None, disabled.
|
||||
shuffle(bool, optional): whether to shuffle indices order before generating
|
||||
batch indices. Default False, don't shuffle indices before generating batch indices.
|
||||
batch_size(int, optional): sample indice number in a mini-batch indices. default 1, each mini-batch includes 1 sample.
|
||||
drop_last(bool, optional): whether drop the last incomplete (less than 1 mini-batch) batch dataset. Default False, keep it.
|
||||
see :ref:`api_paddle_io_DataLoader`
|
||||
|
||||
Returns:
|
||||
BatchSampler: an iterable object for indices iterating
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import RandomSampler, BatchSampler, Dataset
|
||||
|
||||
>>> np.random.seed(2023)
|
||||
>>> # init with dataset
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([784]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> bs = BatchSampler(
|
||||
... dataset=RandomDataset(100),
|
||||
... shuffle=False,
|
||||
... batch_size=16,
|
||||
... drop_last=False,
|
||||
... )
|
||||
>>> for batch_indices in bs:
|
||||
... print(batch_indices)
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
|
||||
...
|
||||
[96, 97, 98, 99]
|
||||
>>> # init with sampler
|
||||
>>> sampler = RandomSampler(RandomDataset(100))
|
||||
>>> bs = BatchSampler(
|
||||
... sampler=sampler,
|
||||
... batch_size=8,
|
||||
... drop_last=True,
|
||||
... )
|
||||
>>> for batch_indices in bs:
|
||||
... print(batch_indices)
|
||||
[56, 12, 68, 0, 82, 66, 91, 44]
|
||||
...
|
||||
[53, 17, 22, 86, 52, 3, 92, 33]
|
||||
"""
|
||||
|
||||
sampler: Sampler[int] | Iterable[int]
|
||||
batch_size: int
|
||||
shuffle: bool
|
||||
drop_last: bool
|
||||
|
||||
@overload
|
||||
def __init__(
|
||||
self,
|
||||
dataset: Sized | None = None,
|
||||
sampler: Sampler | Iterable[int] | None = None,
|
||||
shuffle: bool = False,
|
||||
batch_size: int = 1,
|
||||
drop_last: bool = False,
|
||||
) -> None: ...
|
||||
|
||||
@overload
|
||||
def __init__(
|
||||
self,
|
||||
sampler: Sampler | Iterable[int] | None = None,
|
||||
batch_size: int = 1,
|
||||
drop_last: bool = False,
|
||||
) -> None: ...
|
||||
|
||||
@batch_sampler_decorator
|
||||
def __init__(
|
||||
self,
|
||||
dataset: Sized | None = None,
|
||||
sampler: Sampler | Iterable[int] | None = None,
|
||||
shuffle: bool = False,
|
||||
batch_size: int = 1,
|
||||
drop_last: bool = False,
|
||||
) -> None:
|
||||
if dataset is None:
|
||||
assert sampler is not None, (
|
||||
"either dataset or sampler should be set"
|
||||
)
|
||||
assert isinstance(sampler, (Sampler, Iterable)), (
|
||||
f"sampler should be either paddle.io.Sampler or Iterable, but got {type(sampler)}"
|
||||
)
|
||||
assert not shuffle, "shuffle should be False when sampler is set"
|
||||
self.sampler = sampler
|
||||
else:
|
||||
assert not isinstance(dataset, IterableDataset), (
|
||||
"dataset should not be a paddle.io.IterableDataset"
|
||||
)
|
||||
assert sampler is None, "should not set both dataset and sampler"
|
||||
assert isinstance(shuffle, bool), (
|
||||
f"shuffle should be a boolean value, but got {type(shuffle)}"
|
||||
)
|
||||
if shuffle:
|
||||
self.sampler = RandomSampler(dataset)
|
||||
else:
|
||||
self.sampler = SequenceSampler(dataset)
|
||||
|
||||
assert isinstance(batch_size, int) and batch_size > 0, (
|
||||
f"batch_size should be a positive integer, but got {batch_size}"
|
||||
)
|
||||
self.batch_size = batch_size # per_device_batch_size or mini_batch_size
|
||||
self.shuffle = shuffle
|
||||
assert isinstance(drop_last, bool), (
|
||||
f"drop_last should be a boolean value, but got {type(drop_last)}"
|
||||
)
|
||||
self.drop_last = drop_last
|
||||
|
||||
# TODO(dev): consider to make it as public argument, acc_steps is only used
|
||||
# in auto-parallel
|
||||
self._acc_steps = 1
|
||||
|
||||
def __iter__(self) -> Iterator[list[int]]:
|
||||
local_batch_size = self.batch_size * self._acc_steps
|
||||
batch_indices = []
|
||||
for idx in self.sampler:
|
||||
batch_indices.append(idx)
|
||||
if len(batch_indices) == local_batch_size:
|
||||
yield batch_indices
|
||||
batch_indices = []
|
||||
if not self.drop_last and len(batch_indices) > 0:
|
||||
yield batch_indices
|
||||
|
||||
def __len__(self) -> int:
|
||||
local_batch_size = self.batch_size * self._acc_steps
|
||||
num_samples = len(self.sampler)
|
||||
num_samples += int(not self.drop_last) * (local_batch_size - 1)
|
||||
return num_samples // local_batch_size
|
||||
|
||||
|
||||
class _InfiniteIterableSampler(Sampler[Sequence[None]]):
|
||||
dataset: IterableDataset
|
||||
batch_size: int
|
||||
|
||||
def __init__(self, dataset: IterableDataset, batch_size: int = 1) -> None:
|
||||
assert isinstance(dataset, IterableDataset), (
|
||||
"dataset should be an instance of paddle.io.IterableDataset"
|
||||
)
|
||||
self.dataset = dataset
|
||||
self.batch_size = batch_size
|
||||
|
||||
def __iter__(self) -> Iterator[list[None]]:
|
||||
while True:
|
||||
yield [None] * self.batch_size
|
||||
|
||||
|
||||
class DistributedBatchSampler(BatchSampler):
|
||||
"""Sampler that restricts data loading to a subset of the dataset.
|
||||
|
||||
In such case, each process can pass a DistributedBatchSampler instance
|
||||
as a DataLoader sampler, and load a subset of the original dataset that
|
||||
is exclusive to it.
|
||||
|
||||
.. note::
|
||||
Dataset is assumed to be of constant size.
|
||||
|
||||
Args:
|
||||
dataset(Dataset): this could be an instance of subclass of :ref:`api_paddle_io_Dataset`
|
||||
or other python object which implemented
|
||||
`__len__` for BatchSampler to get indices of samples.
|
||||
batch_size(int): sample size of each mini-batch.
|
||||
num_replicas(int, optional): process number in distributed training.
|
||||
If :attr:`num_replicas` is None, :attr:`num_replicas` will be
|
||||
retrieved from :ref:`api_paddle_distributed_ParallelEnv` .
|
||||
Default None.
|
||||
rank(int, optional): the rank of the current process among :attr:`num_replicas`
|
||||
processes. If :attr:`rank` is None, :attr:`rank` is retrieved from
|
||||
:ref:`api_paddle_distributed_ParallelEnv`. Default None.
|
||||
shuffle(bool, optional): whether to shuffle indices order before generating
|
||||
batch indices. Default False.
|
||||
drop_last(bool, optional): whether drop the last incomplete(less than a mini-batch) batch dataset size.
|
||||
Default False.
|
||||
seed(int, optional): random seed used to shuffle indices order if
|
||||
:attr:`shuffle=True`. This number should be identical across all
|
||||
processes in the distributed group. Default 0.
|
||||
|
||||
Returns:
|
||||
DistributedBatchSampler, return an iterable object for indices iterating.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
|
||||
>>> from paddle.io import Dataset, DistributedBatchSampler
|
||||
|
||||
>>> # init with dataset
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([784]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> dataset = RandomDataset(100)
|
||||
>>> sampler = DistributedBatchSampler(dataset, batch_size=64)
|
||||
|
||||
>>> for data in sampler:
|
||||
... # do something
|
||||
... break
|
||||
"""
|
||||
|
||||
dataset: Sized
|
||||
batch_size: int
|
||||
drop_last: bool
|
||||
nranks: int
|
||||
epoch: int
|
||||
local_rank: int
|
||||
num_samples: int
|
||||
total_size: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset: Sized,
|
||||
batch_size: int,
|
||||
num_replicas: int | None = None,
|
||||
rank: int | None = None,
|
||||
shuffle: bool = False,
|
||||
drop_last: bool = False,
|
||||
seed: int = 0,
|
||||
) -> None:
|
||||
self.dataset = dataset
|
||||
|
||||
assert isinstance(batch_size, int) and batch_size > 0, (
|
||||
"batch_size should be a positive integer"
|
||||
)
|
||||
self.batch_size = batch_size
|
||||
assert isinstance(shuffle, bool), "shuffle should be a boolean value"
|
||||
self.shuffle = shuffle
|
||||
assert isinstance(drop_last, bool), (
|
||||
"drop_last should be a boolean number"
|
||||
)
|
||||
|
||||
from paddle.distributed import ParallelEnv
|
||||
|
||||
if num_replicas is not None:
|
||||
assert isinstance(num_replicas, int) and num_replicas > 0, (
|
||||
"num_replicas should be a positive integer"
|
||||
)
|
||||
self.nranks = num_replicas
|
||||
else:
|
||||
self.nranks = ParallelEnv().nranks
|
||||
|
||||
if rank is not None:
|
||||
assert isinstance(rank, int) and rank >= 0, (
|
||||
"rank should be a non-negative integer"
|
||||
)
|
||||
self.local_rank = rank
|
||||
else:
|
||||
self.local_rank = ParallelEnv().local_rank
|
||||
|
||||
self.drop_last = drop_last
|
||||
self.epoch = 0
|
||||
self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.nranks))
|
||||
self.total_size = self.num_samples * self.nranks
|
||||
self.seed = seed
|
||||
|
||||
# TODO(dev): consider to make it as public argument, acc_steps is only used
|
||||
# in auto-parallel
|
||||
self._acc_steps = 1
|
||||
|
||||
def __iter__(self) -> Iterator[list[int]]:
|
||||
local_batch_size = self.batch_size * self._acc_steps
|
||||
num_samples = len(self.dataset)
|
||||
indices = np.arange(num_samples).tolist()
|
||||
# add extra samples to make it evenly divisible
|
||||
padding_size = self.total_size - len(indices)
|
||||
if padding_size <= len(indices):
|
||||
indices += indices[:padding_size]
|
||||
else:
|
||||
indices += (indices * math.ceil(padding_size / len(indices)))[
|
||||
:padding_size
|
||||
]
|
||||
|
||||
assert len(indices) == self.total_size
|
||||
if self.shuffle:
|
||||
np.random.RandomState(self.seed + self.epoch).shuffle(indices)
|
||||
self.epoch += 1
|
||||
|
||||
# subsample
|
||||
def _get_indices_by_batch_size(indices):
|
||||
subsampled_indices = []
|
||||
last_batch_size = self.total_size % (self.batch_size * self.nranks)
|
||||
assert last_batch_size % self.nranks == 0
|
||||
last_local_batch_size = last_batch_size // self.nranks
|
||||
|
||||
for i in range(
|
||||
self.local_rank * self.batch_size,
|
||||
len(indices) - last_batch_size,
|
||||
self.batch_size * self.nranks,
|
||||
):
|
||||
subsampled_indices.extend(indices[i : i + self.batch_size])
|
||||
|
||||
indices = indices[len(indices) - last_batch_size :]
|
||||
subsampled_indices.extend(
|
||||
indices[
|
||||
self.local_rank * last_local_batch_size : (
|
||||
self.local_rank + 1
|
||||
)
|
||||
* last_local_batch_size
|
||||
]
|
||||
)
|
||||
|
||||
return subsampled_indices
|
||||
|
||||
if self.nranks > 1:
|
||||
indices = _get_indices_by_batch_size(indices)
|
||||
|
||||
assert len(indices) == self.num_samples
|
||||
_sample_iter = iter(indices)
|
||||
|
||||
batch_indices = []
|
||||
for idx in _sample_iter:
|
||||
batch_indices.append(idx)
|
||||
if len(batch_indices) == local_batch_size:
|
||||
yield batch_indices
|
||||
batch_indices = []
|
||||
if not self.drop_last and len(batch_indices) > 0:
|
||||
yield batch_indices
|
||||
|
||||
def __len__(self) -> int:
|
||||
local_batch_size = self.batch_size * self._acc_steps
|
||||
num_samples = self.num_samples
|
||||
num_samples += int(not self.drop_last) * (local_batch_size - 1)
|
||||
return num_samples // local_batch_size
|
||||
|
||||
def set_epoch(self, epoch: int) -> None:
|
||||
"""
|
||||
Sets the epoch number. When :attr:`shuffle=True`, this number is used
|
||||
as seeds of random numbers. By default, users may not set this, all
|
||||
replicas (workers) use a different random ordering for each epoch.
|
||||
If set same number at each epoch, this sampler will yield the same
|
||||
ordering at all epochs.
|
||||
|
||||
Arguments:
|
||||
epoch (int): Epoch number.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
|
||||
>>> from paddle.io import Dataset, DistributedBatchSampler
|
||||
|
||||
>>> # init with dataset
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([784]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> dataset = RandomDataset(100)
|
||||
>>> sampler = DistributedBatchSampler(dataset, batch_size=64)
|
||||
|
||||
>>> for epoch in range(10):
|
||||
... sampler.set_epoch(epoch)
|
||||
"""
|
||||
self.epoch = epoch
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import numbers
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
|
||||
from ...framework import core
|
||||
|
||||
|
||||
def default_collate_fn(batch):
|
||||
"""
|
||||
Default batch collating function for :code:`paddle.io.DataLoader`,
|
||||
get input data as a list of sample datas, each element in list
|
||||
if the data of a sample, and sample data should composed of list,
|
||||
dictionary, string, number, numpy array and paddle.Tensor, this
|
||||
function will parse input data recursively and stack number,
|
||||
numpy array and paddle.Tensor datas as batch datas. e.g. for
|
||||
following input data:
|
||||
|
||||
[{'image': np.array(shape=[3, 224, 224]), 'label': 1},
|
||||
{'image': np.array(shape=[3, 224, 224]), 'label': 3},
|
||||
{'image': np.array(shape=[3, 224, 224]), 'label': 4},
|
||||
{'image': np.array(shape=[3, 224, 224]), 'label': 5},]
|
||||
|
||||
|
||||
This default collate function zipped each number and numpy array
|
||||
field together and stack each field as the batch field as follows:
|
||||
|
||||
{'image': np.array(shape=[4, 3, 224, 224]), 'label': np.array([1, 3, 4, 5])}
|
||||
|
||||
|
||||
Args:
|
||||
batch(list of sample data): batch should be a list of sample data.
|
||||
|
||||
Returns:
|
||||
Batched data: batched each number, numpy array and paddle.Tensor
|
||||
in input data.
|
||||
"""
|
||||
sample = batch[0]
|
||||
if isinstance(sample, np.ndarray):
|
||||
batch = np.stack(batch, axis=0)
|
||||
return batch
|
||||
elif isinstance(sample, paddle.Tensor):
|
||||
return paddle.stack(batch, axis=0)
|
||||
elif isinstance(sample, numbers.Number):
|
||||
batch = np.array(batch)
|
||||
return batch
|
||||
elif isinstance(sample, (str, bytes)):
|
||||
return batch
|
||||
elif isinstance(sample, Mapping):
|
||||
return {
|
||||
key: default_collate_fn([d[key] for d in batch]) for key in sample
|
||||
}
|
||||
elif isinstance(sample, Sequence):
|
||||
sample_fields_num = len(sample)
|
||||
if not all(len(sample) == sample_fields_num for sample in iter(batch)):
|
||||
raise RuntimeError(
|
||||
"fields number not same among samples in a batch"
|
||||
)
|
||||
return [default_collate_fn(fields) for fields in zip(*batch)]
|
||||
|
||||
raise TypeError(
|
||||
"batch data con only contains: tensor, numpy.ndarray, "
|
||||
f"dict, list, number, but got {type(sample)}"
|
||||
)
|
||||
|
||||
|
||||
def default_convert_fn(batch):
|
||||
"""
|
||||
Default batch converting function for :code:`paddle.io.DataLoader`.
|
||||
get input data as a list of sample datas, each element in list
|
||||
if the data of a sample, and sample data should composed of list,
|
||||
dictionary, string, number, numpy array and paddle.Tensor.
|
||||
|
||||
.. note::
|
||||
This function is default :attr:`collate_fn` in **Disable
|
||||
automatic batching** mode, for **Disable automatic batching**
|
||||
mode, please ses :attr:`paddle.io.DataLoader`
|
||||
|
||||
Args:
|
||||
batch(list of sample data): batch should be a list of sample data.
|
||||
|
||||
Returns:
|
||||
Batched data: batched each number, numpy array and paddle.Tensor
|
||||
in input data.
|
||||
"""
|
||||
if isinstance(batch, (paddle.Tensor, np.ndarray, core.eager.Tensor)):
|
||||
return batch
|
||||
elif isinstance(batch, (str, bytes)):
|
||||
return batch
|
||||
elif isinstance(batch, Mapping):
|
||||
return {key: default_convert_fn(batch[key]) for key in batch}
|
||||
elif isinstance(batch, Sequence):
|
||||
return [default_convert_fn(d) for d in batch]
|
||||
else:
|
||||
return batch
|
||||
@@ -0,0 +1,887 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import itertools
|
||||
import logging
|
||||
import os
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle import profiler
|
||||
from paddle.base.framework import _current_expected_place, _set_expected_place
|
||||
from paddle.pir.core import datatype_to_vartype
|
||||
from paddle.profiler.timer import benchmark
|
||||
from paddle.profiler.utils import in_profiler_mode
|
||||
|
||||
from ...framework import core, in_dynamic_mode, in_pir_mode
|
||||
from ..multiprocess_utils import (
|
||||
MP_STATUS_CHECK_INTERVAL,
|
||||
CleanupFuncRegistrar,
|
||||
_set_SIGCHLD_handler,
|
||||
)
|
||||
from .batch_sampler import _InfiniteIterableSampler
|
||||
from .collate import default_collate_fn, default_convert_fn
|
||||
from .flat import _flatten_batch, _restore_batch
|
||||
from .worker import (
|
||||
_DatasetKind,
|
||||
_IterableDatasetStopIteration,
|
||||
_ResumeIteration,
|
||||
_worker_loop,
|
||||
_WorkerException,
|
||||
)
|
||||
|
||||
# NOTE: fix `terminate called without an active exception`
|
||||
# if for loop break and program exit immediately(with no model
|
||||
# layers processing) after iterate **the first few data** in
|
||||
# distributed launch mode, distributed launch will call
|
||||
# terminate() to kill main process on each devices, but thread
|
||||
# is still iterating to fulfill blocking queue caches, which
|
||||
# may cause thread error `terminate called without an active
|
||||
# exception` for terminate is a strong signal and `__del__`
|
||||
# of DataLoader may not be called, so we add a global link to
|
||||
# the last DataLoader instance to call `__del__` to clean up
|
||||
# resources
|
||||
# NOTE: cannot simply as `__del__` to CleanupFuncRegistrar,
|
||||
# for this will remain a link to each DataLoader instance in
|
||||
# global, and will precludes GC to auto collect DataLoader
|
||||
# instance and will cause memory leak
|
||||
_loader = None
|
||||
|
||||
|
||||
def _clear_loader():
|
||||
global _loader
|
||||
if _loader is not None:
|
||||
try:
|
||||
_loader.__del__()
|
||||
del _loader
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
CleanupFuncRegistrar.register(_clear_loader)
|
||||
|
||||
|
||||
class _DataLoaderIterBase:
|
||||
"""
|
||||
Iterator implement of DataLoader, will load and feed mini-batch
|
||||
data by setting in given dataloader.
|
||||
|
||||
Args:
|
||||
loader(instance of DataLoader): instance of `paddle.io.DataLoader`
|
||||
"""
|
||||
|
||||
def __init__(self, loader):
|
||||
self._dataset = loader.dataset
|
||||
self._feed_list = loader.feed_list or []
|
||||
self._places = loader.places
|
||||
self._return_list = loader.return_list
|
||||
self._batch_sampler = loader.batch_sampler
|
||||
self._drop_last = loader.drop_last
|
||||
self._auto_collate_batch = loader.auto_collate_batch
|
||||
self._num_workers = loader.num_workers
|
||||
self._use_buffer_reader = loader.use_buffer_reader
|
||||
self._reader_buffer_size = loader.reader_buffer_size
|
||||
self._prefetch_factor = loader.prefetch_factor
|
||||
self._use_shared_memory = loader.use_shared_memory
|
||||
self._timeout = (
|
||||
loader.timeout if loader.timeout > 0 else MP_STATUS_CHECK_INTERVAL
|
||||
)
|
||||
self._worker_init_fn = loader.worker_init_fn
|
||||
self._dataset_kind = loader.dataset_kind
|
||||
self._pin_memory = loader.pin_memory
|
||||
|
||||
self._sampler_iter = iter(self._index_sampler)
|
||||
if self._auto_collate_batch:
|
||||
self._collate_fn = loader.collate_fn or default_collate_fn
|
||||
else:
|
||||
self._collate_fn = loader.collate_fn or default_convert_fn
|
||||
|
||||
# DenseTensorBlockingQueue instance for create_py_reader and a thread
|
||||
# to put mini-batch data to self._blocking_queue, mini-batch data
|
||||
# will be get from:
|
||||
# 1. multi-process mode: get data from workers' result queue
|
||||
# 2. single-process mode: read mini-batch data in main process
|
||||
self._blocking_queue = None
|
||||
self._thread = None
|
||||
self._thread_done_event = threading.Event()
|
||||
|
||||
@property
|
||||
def _index_sampler(self):
|
||||
if self._auto_collate_batch:
|
||||
return self._batch_sampler
|
||||
else:
|
||||
if self._dataset_kind == _DatasetKind.MAP:
|
||||
return list(range(len(self._dataset)))
|
||||
else:
|
||||
return _InfiniteIterableSampler(self._dataset, 1)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
raise NotImplementedError('Should implement `__next__` for a iterator')
|
||||
|
||||
def __len__(self):
|
||||
return len(self._batch_sampler)
|
||||
|
||||
def _exit_thread_expectedly(self):
|
||||
self._thread_done_event.set()
|
||||
if self._blocking_queue:
|
||||
self._blocking_queue.close()
|
||||
|
||||
def _exit_thread_unexpectedly(self):
|
||||
self._thread_done_event.set()
|
||||
if self._blocking_queue:
|
||||
self._blocking_queue.kill()
|
||||
|
||||
|
||||
class _DataLoaderIterSingleProcess(_DataLoaderIterBase):
|
||||
"""
|
||||
Single process implement of DataLoaderIter, loading data from
|
||||
loader.data in main process
|
||||
"""
|
||||
|
||||
def __init__(self, loader):
|
||||
super().__init__(loader)
|
||||
|
||||
self._dataset_fetcher = _DatasetKind.create_fetcher(
|
||||
self._dataset_kind,
|
||||
self._dataset,
|
||||
self._auto_collate_batch,
|
||||
self._collate_fn,
|
||||
self._drop_last,
|
||||
)
|
||||
|
||||
# NOTE: _structure_infos used to record the data structure of
|
||||
# batch to restore batch structure after reading Tensor
|
||||
# from blocking_queue in single-process mode. Note that
|
||||
# only single process is used in single-process mode, we
|
||||
# can record the data structure sequencely in a list without
|
||||
# recording the send and recv index
|
||||
self._structure_infos = []
|
||||
|
||||
# NOTE: len(self._places) batch data compose as an output
|
||||
# iteration, set blocking_queue can cache "self._prefetch_factor" iteration datas
|
||||
# at most here
|
||||
self._blocking_queue_capacity = self._prefetch_factor * len(
|
||||
self._places
|
||||
)
|
||||
|
||||
self._shutdown = False
|
||||
try:
|
||||
self._init_thread()
|
||||
except Exception:
|
||||
self._try_shutdown_all()
|
||||
raise
|
||||
|
||||
global _loader
|
||||
_loader = self
|
||||
|
||||
def _init_thread(self):
|
||||
self._var_names = [v.name for v in self._feed_list]
|
||||
self._shapes = [v.shape for v in self._feed_list]
|
||||
if in_pir_mode():
|
||||
self._need_check_feed = [False for v in self._feed_list]
|
||||
self._dtypes = [
|
||||
datatype_to_vartype[v.dtype] for v in self._feed_list
|
||||
]
|
||||
else:
|
||||
self._need_check_feed = [
|
||||
v.desc.need_check_feed() for v in self._feed_list
|
||||
]
|
||||
self._dtypes = [v.dtype for v in self._feed_list]
|
||||
# if only 1 place, do not need to keep order
|
||||
self._blocking_queue = core.init_dense_tensor_blocking_queue(
|
||||
core.Variable(),
|
||||
self._blocking_queue_capacity,
|
||||
len(self._places) > 1,
|
||||
)
|
||||
self._reader = core.create_py_reader(
|
||||
self._blocking_queue,
|
||||
self._var_names,
|
||||
self._shapes,
|
||||
self._dtypes,
|
||||
self._need_check_feed,
|
||||
self._places,
|
||||
self._use_buffer_reader,
|
||||
True,
|
||||
self._pin_memory,
|
||||
self._reader_buffer_size,
|
||||
)
|
||||
|
||||
self._thread = threading.Thread(
|
||||
target=self._thread_loop, args=(_current_expected_place(),)
|
||||
)
|
||||
self._thread.daemon = True
|
||||
self._thread.start()
|
||||
|
||||
def _thread_loop(self, legacy_expected_place):
|
||||
# NOTE(zhiqiu): Set the expected place for new thread as the same as father thread,
|
||||
# and it will call platform::SetDeviceId() in c++ internally.
|
||||
# If we do not set cudaDeviceId in new thread, the default cudaDeviceId will be 0,
|
||||
# Which may cost hundreds of MB of GPU memory on CUDAPlace(0) if calling some cuda
|
||||
# APIs in this thread.
|
||||
core.set_current_thread_name("Dataloader_" + str(id(self)))
|
||||
_set_expected_place(legacy_expected_place)
|
||||
|
||||
while not self._thread_done_event.is_set():
|
||||
try:
|
||||
indices = next(self._sampler_iter)
|
||||
|
||||
# read data from dataset in mini-batch
|
||||
# with paddle.base.dygraph.guard(place=paddle.CPUPlace()):
|
||||
# read data from dataset in mini-batch
|
||||
batch = self._dataset_fetcher.fetch(
|
||||
indices, self._thread_done_event
|
||||
)
|
||||
except StopIteration:
|
||||
self._exit_thread_expectedly()
|
||||
return
|
||||
|
||||
if batch is None or self._thread_done_event.is_set():
|
||||
break
|
||||
|
||||
# flat batch and record structure infos
|
||||
batch, structure = _flatten_batch(batch)
|
||||
self._structure_infos.append(structure)
|
||||
|
||||
if self._thread_done_event.is_set():
|
||||
break
|
||||
|
||||
try:
|
||||
# pack as DenseTensorArray
|
||||
array = core.DenseTensorArray()
|
||||
for slot in batch:
|
||||
if isinstance(slot, paddle.Tensor):
|
||||
slot = slot.value().get_tensor()
|
||||
elif not isinstance(slot, core.DenseTensor):
|
||||
tmp = core.DenseTensor()
|
||||
tmp.set(slot, core.CPUPlace())
|
||||
slot = tmp
|
||||
|
||||
array.append(slot)
|
||||
|
||||
if self._thread_done_event.is_set():
|
||||
break
|
||||
|
||||
try:
|
||||
self._blocking_queue.push(array)
|
||||
except:
|
||||
self._exit_thread_expectedly()
|
||||
|
||||
except Exception as e:
|
||||
self._exit_thread_unexpectedly()
|
||||
raise e
|
||||
|
||||
self._exit_thread_expectedly()
|
||||
|
||||
def __next__(self):
|
||||
if in_profiler_mode():
|
||||
trace_event = profiler.RecordEvent(
|
||||
name="_DataLoaderIterSingleProcess",
|
||||
event_type=profiler.TracerEventType.Dataloader,
|
||||
)
|
||||
trace_event.begin()
|
||||
try:
|
||||
benchmark().check_if_need_record(self)
|
||||
benchmark().before_reader()
|
||||
if in_dynamic_mode():
|
||||
data = core.eager.read_next_tensor_list(
|
||||
self._reader.read_next_list()[0]
|
||||
)
|
||||
data = _restore_batch(data, self._structure_infos.pop(0))
|
||||
else:
|
||||
# in static graph mode
|
||||
if self._return_list:
|
||||
data = self._reader.read_next_list()
|
||||
for i in range(len(data)):
|
||||
data[i] = data[i]._move_to_list()
|
||||
structs = [
|
||||
self._structure_infos.pop(0)
|
||||
for _ in range(len(self._places))
|
||||
]
|
||||
data = [_restore_batch(d, s) for d, s in zip(data, structs)]
|
||||
# static graph organized data on multi-device with list, if
|
||||
# place number is 1, there is only 1 device, extra the data
|
||||
# from list for devices to be compatible with dygraph mode
|
||||
if len(self._places) == 1:
|
||||
data = data[0]
|
||||
else:
|
||||
data = self._reader.read_next()
|
||||
benchmark().after_reader()
|
||||
|
||||
return data
|
||||
except StopIteration:
|
||||
self._reader.shutdown()
|
||||
self._try_shutdown_all()
|
||||
raise
|
||||
finally:
|
||||
if in_profiler_mode():
|
||||
trace_event.end()
|
||||
|
||||
def _shutdown_thread(self):
|
||||
if self._thread:
|
||||
self._thread_done_event.set()
|
||||
# NOTE: we wait for _thread exit for 3 seconds, if
|
||||
# thread not exit normally, force kill it
|
||||
for _ in range(3):
|
||||
if self._thread.is_alive():
|
||||
time.sleep(1)
|
||||
else:
|
||||
break
|
||||
else:
|
||||
if self._thread is not threading.current_thread():
|
||||
self._thread.join()
|
||||
|
||||
self._thread = None
|
||||
|
||||
def _try_shutdown_all(self):
|
||||
if not self._shutdown:
|
||||
try:
|
||||
# # _blocking_queue in keep order mode holds sub-threads
|
||||
# # need to release thread resources on unexpected exit
|
||||
if self._blocking_queue:
|
||||
self._blocking_queue.close()
|
||||
self._blocking_queue = None
|
||||
# NOTE: blocking queue should be closed firstly for
|
||||
# blocking queue read may hang and _thread_done_event
|
||||
# cannot be checked
|
||||
self._shutdown_thread()
|
||||
finally:
|
||||
self._shutdown = True
|
||||
|
||||
def __del__(self):
|
||||
self._try_shutdown_all()
|
||||
|
||||
|
||||
class _DataLoaderIterMultiProcess(_DataLoaderIterBase):
|
||||
def __init__(self, loader):
|
||||
super().__init__(loader)
|
||||
|
||||
self._persistent_workers = loader._persistent_workers
|
||||
self._resume_worker_cnt = 0
|
||||
|
||||
assert self._num_workers > 0, (
|
||||
f"Multi-process DataLoader invalid num_workers({self._num_workers})"
|
||||
)
|
||||
|
||||
# subprocess workers' result queue
|
||||
self._data_queue = None
|
||||
|
||||
# data get from _data_queue will be reordered by _rcvd_idx
|
||||
# for data order keeping, data index not equal _rcvd_idx
|
||||
# will be cached in _task_infos
|
||||
self._send_idx = 0
|
||||
self._rcvd_idx = 0
|
||||
self._batches_outstanding = 0
|
||||
self._task_infos = {}
|
||||
self._structure_infos = []
|
||||
|
||||
# indices outstand as _outstanding_capacity at first, and
|
||||
# blocking_queue capacity is also _outstanding_capacity.
|
||||
# _outstanding_capacity here to make sure each indices_queue
|
||||
# has at least "_prefetch_factor" indices, and outstanding batch cached
|
||||
# output data for at least "_prefetch_factor" iterations(Note that len(_places)
|
||||
# batches will be composed as an iteration output)
|
||||
self._outstanding_capacity = self._prefetch_factor * max(
|
||||
self._num_workers, len(self._places)
|
||||
)
|
||||
|
||||
# see _try_put_indices
|
||||
self._thread_lock = threading.Lock()
|
||||
|
||||
self._base_seed = np.random.randint(low=0, high=sys.maxsize)
|
||||
|
||||
# Note(zhangbo): shm_buffer_size is used for MemoryMapAllocationPool.
|
||||
# MemoryMapAllocationPool is used to cache and reuse shm, thus reducing munmap in dataloader.
|
||||
# For more details, please see: paddle/base/memory/allocation/mmap_allocator.h
|
||||
if os.environ.get('FLAGS_use_shm_cache', False) in [
|
||||
1,
|
||||
'1',
|
||||
True,
|
||||
'True',
|
||||
'true',
|
||||
]:
|
||||
try:
|
||||
self._worker_shm_buffer_size = (2 + 1) * len(self._dataset[0])
|
||||
except:
|
||||
self._worker_shm_buffer_size = 0
|
||||
warnings.warn(
|
||||
"Setting the shm cache buffer size to 0, equivalent to not using the shm cache policy."
|
||||
)
|
||||
else:
|
||||
self._worker_shm_buffer_size = 0
|
||||
self._main_thread_shm_buffer_size = (
|
||||
(self._worker_shm_buffer_size) * 2 * self._num_workers
|
||||
)
|
||||
|
||||
self._shutdown = False
|
||||
# init workers and indices queues and put 2 indices in each indices queue
|
||||
self._init_workers()
|
||||
for _ in range(self._outstanding_capacity):
|
||||
self._try_put_indices()
|
||||
|
||||
try:
|
||||
self._init_thread()
|
||||
except Exception:
|
||||
self._try_shutdown_all()
|
||||
raise
|
||||
|
||||
def _init_workers(self):
|
||||
from paddle.incubate import multiprocessing
|
||||
|
||||
# multiprocess worker and indice queue list initial as empty
|
||||
self._workers = []
|
||||
self._worker_status = []
|
||||
self._indices_queues = []
|
||||
self._workers_idx_cycle = itertools.cycle(range(self._num_workers))
|
||||
|
||||
# create data_queue for workers
|
||||
self._data_queue = multiprocessing.Queue()
|
||||
|
||||
# event for workers and thread, thread event is only need
|
||||
# in multi-processing mode
|
||||
self._workers_done_event = multiprocessing.Event()
|
||||
self._thread_done_event = threading.Event()
|
||||
|
||||
for i in range(self._num_workers):
|
||||
indices_queue = multiprocessing.Queue()
|
||||
indices_queue.cancel_join_thread()
|
||||
self._indices_queues.append(indices_queue)
|
||||
worker = multiprocessing.Process(
|
||||
target=_worker_loop,
|
||||
args=(
|
||||
self._dataset,
|
||||
self._dataset_kind,
|
||||
indices_queue,
|
||||
self._data_queue,
|
||||
self._workers_done_event,
|
||||
self._auto_collate_batch,
|
||||
self._collate_fn,
|
||||
self._drop_last,
|
||||
self._worker_init_fn,
|
||||
i,
|
||||
self._num_workers,
|
||||
self._use_shared_memory,
|
||||
self._base_seed,
|
||||
self._worker_shm_buffer_size,
|
||||
),
|
||||
)
|
||||
worker.daemon = True
|
||||
worker.start()
|
||||
self._workers.append(worker)
|
||||
self._worker_status.append(True)
|
||||
|
||||
core._set_process_pids(id(self), tuple(w.pid for w in self._workers))
|
||||
_set_SIGCHLD_handler()
|
||||
|
||||
def _clear_and_remove_data_queue(self):
|
||||
if self._data_queue is not None:
|
||||
while True:
|
||||
try:
|
||||
self._data_queue.get_nowait()
|
||||
except:
|
||||
self._data_queue.cancel_join_thread()
|
||||
self._data_queue.close()
|
||||
break
|
||||
|
||||
def _init_thread(self):
|
||||
self._var_names = [v.name for v in self._feed_list]
|
||||
self._shapes = [v.shape for v in self._feed_list]
|
||||
if in_pir_mode():
|
||||
self._need_check_feed = [False for v in self._feed_list]
|
||||
self._dtypes = [
|
||||
datatype_to_vartype[v.dtype] for v in self._feed_list
|
||||
]
|
||||
else:
|
||||
self._need_check_feed = [
|
||||
v.desc.need_check_feed() for v in self._feed_list
|
||||
]
|
||||
self._dtypes = [v.dtype for v in self._feed_list]
|
||||
# if only 1 place, do not need to keep order
|
||||
self._blocking_queue = core.init_dense_tensor_blocking_queue(
|
||||
core.Variable(), self._outstanding_capacity, len(self._places) > 1
|
||||
)
|
||||
core._set_max_memory_map_allocation_pool_size(
|
||||
self._main_thread_shm_buffer_size
|
||||
)
|
||||
self._reader = core.create_py_reader(
|
||||
self._blocking_queue,
|
||||
self._var_names,
|
||||
self._shapes,
|
||||
self._dtypes,
|
||||
self._need_check_feed,
|
||||
self._places,
|
||||
self._use_buffer_reader,
|
||||
True,
|
||||
self._pin_memory,
|
||||
self._reader_buffer_size,
|
||||
)
|
||||
|
||||
self._thread_done_event = threading.Event()
|
||||
# thread event is only need in multi-processing mode
|
||||
self._thread = threading.Thread(
|
||||
target=self._thread_loop, args=(_current_expected_place(),)
|
||||
)
|
||||
self._thread.daemon = True
|
||||
self._thread.start()
|
||||
|
||||
def _reset(self):
|
||||
# resume iteration in following steps
|
||||
# 1. Resume workers, clear worker caches
|
||||
# put _ResumeIteration to all worker as resume iteration flag
|
||||
with self._thread_lock:
|
||||
self._resume_worker_cnt = self._num_workers
|
||||
for worker_id in range(self._num_workers):
|
||||
self._indices_queues[worker_id].put(_ResumeIteration())
|
||||
self._batches_outstanding += 1
|
||||
# all flag will be check in _thread_loop, simply wait here
|
||||
while self._resume_worker_cnt > 0:
|
||||
time.sleep(0.5)
|
||||
|
||||
# 2. clear blocking_queue caches
|
||||
# in order not to restart the thread, we just clear
|
||||
# the blocking_queue cachees instead of recreating one
|
||||
while self._blocking_queue.size() >= len(self._places):
|
||||
if in_dynamic_mode():
|
||||
data = core.eager.read_next_tensor_list(
|
||||
self._reader.read_next_list()[0]
|
||||
)
|
||||
else:
|
||||
if self._return_list:
|
||||
self._reader.read_next_list()
|
||||
else:
|
||||
data = self._reader.read_next()
|
||||
|
||||
# 3. reset all states
|
||||
self._send_idx = 0
|
||||
self._rcvd_idx = 0
|
||||
self._batches_outstanding = 0
|
||||
self._task_infos = {}
|
||||
self._structure_infos = []
|
||||
|
||||
# set all worker status available
|
||||
self._worker_status = [True] * self._num_workers
|
||||
|
||||
# 4. reset _sampler_iter and put prefetch indices to start next epoch
|
||||
# init workers and indices queues and put 2 indices in each indices queue
|
||||
self._sampler_iter = iter(self._index_sampler)
|
||||
for _ in range(self._outstanding_capacity):
|
||||
self._try_put_indices()
|
||||
|
||||
def _shutdown_worker(self, worker_id, shutdown=False):
|
||||
if worker_id < len(self._worker_status) and (
|
||||
self._worker_status[worker_id]
|
||||
or self._persistent_workers
|
||||
and shutdown
|
||||
):
|
||||
self._indices_queues[worker_id].put(None)
|
||||
self._worker_status[worker_id] = False
|
||||
|
||||
def _try_shutdown_all(self, timeout=None):
|
||||
if not self._shutdown:
|
||||
try:
|
||||
self._exit_thread_expectedly()
|
||||
self._clear_and_remove_data_queue()
|
||||
|
||||
# set _workers_done_event should be set before put None
|
||||
# to indices_queue, workers will exit on reading None from
|
||||
# indices_queue
|
||||
self._workers_done_event.set()
|
||||
for i in range(self._num_workers):
|
||||
self._shutdown_worker(i, shutdown=True)
|
||||
|
||||
if not self._shutdown:
|
||||
for w in self._workers:
|
||||
w.join(timeout)
|
||||
for q in self._indices_queues:
|
||||
q.cancel_join_thread()
|
||||
q.close()
|
||||
finally:
|
||||
core._erase_process_pids(id(self))
|
||||
self._shutdown = True
|
||||
|
||||
def _thread_loop(self, legacy_expected_place):
|
||||
# NOTE(zhiqiu): Set the expected place for new thread as the same as father thread,
|
||||
# and it will call platform::SetDeviceId() in c++ internally.
|
||||
# If we do not set cudaDeviceId in new thread, the default cudaDeviceId will be 0,
|
||||
# Which may cost hundreds of MB of GPU memory on CUDAPlace(0) if calling some cuda
|
||||
# APIs in this thread.
|
||||
core.set_current_thread_name("Dataloader_" + str(id(self)))
|
||||
_set_expected_place(legacy_expected_place)
|
||||
|
||||
while not self._thread_done_event.is_set():
|
||||
batch = self._get_data()
|
||||
if not self._thread_done_event.is_set():
|
||||
if batch is None:
|
||||
self._exit_thread_expectedly()
|
||||
else:
|
||||
if isinstance(batch, _ResumeIteration):
|
||||
assert self._resume_worker_cnt > 0
|
||||
self._resume_worker_cnt -= 1
|
||||
continue
|
||||
try:
|
||||
# pack as DenseTensorArray
|
||||
array = core.DenseTensorArray()
|
||||
if self._use_shared_memory:
|
||||
for tensor in batch:
|
||||
array.append(tensor)
|
||||
else:
|
||||
# DenseTensor not in shared memory is not
|
||||
# serializable, cannot be create in workers
|
||||
for slot in batch:
|
||||
if isinstance(slot, paddle.Tensor):
|
||||
slot = slot.get_tensor()
|
||||
elif not isinstance(slot, core.DenseTensor):
|
||||
tmp = core.DenseTensor()
|
||||
tmp.set(slot, core.CPUPlace())
|
||||
slot = tmp
|
||||
array.append(slot)
|
||||
|
||||
if not self._blocking_queue.push(array):
|
||||
self._blocking_queue.close()
|
||||
except Exception as e:
|
||||
self._exit_thread_unexpectedly()
|
||||
raise e
|
||||
finally:
|
||||
self._rcvd_idx += 1
|
||||
|
||||
def _get_data(self):
|
||||
while not self._thread_done_event.is_set():
|
||||
# For IterableDataset, batch indices is generated infinitely
|
||||
# for each worker to raise StopIteration, but a StopIteration
|
||||
# raising process will discard a batch indices which is count
|
||||
# in _send_idx but will not increase _rcvd_idx, so we check
|
||||
# whether the worker is still alive here to skip the discarded
|
||||
# batch indices and increase _rcvd_idx
|
||||
if self._dataset_kind == _DatasetKind.ITER:
|
||||
while self._rcvd_idx < self._send_idx:
|
||||
info = self._task_infos[self._rcvd_idx]
|
||||
if len(info) == 3 or self._worker_status[info[0]]:
|
||||
break
|
||||
del self._task_infos[self._rcvd_idx]
|
||||
self._rcvd_idx += 1
|
||||
self._batches_outstanding -= 1
|
||||
else:
|
||||
# NOTE: when _rcvd_idx catch up _send_idx, which means
|
||||
# one of following:
|
||||
# 1. all 2 * num_workers batches have been loaded
|
||||
# and stored in _blocking_queue
|
||||
# 2. all data drained
|
||||
# we need to let _thread blocking at _data_queue
|
||||
# get_data to inoccupy CPU, otherwise may occupy
|
||||
# CPU time for model running
|
||||
# NOTE: in persistent workers mode, do not check data
|
||||
# drained here, simply let it go to _data_queue
|
||||
# reading to get _ResumeIteration
|
||||
if not self._persistent_workers:
|
||||
# NOTE: _rcvd_idx and _send_idx only record batches among
|
||||
# workers, if batches among workers drained, there
|
||||
# may also be data in blocking queue
|
||||
if self._batches_outstanding < len(self._places):
|
||||
return None
|
||||
|
||||
if (
|
||||
self._rcvd_idx in self._task_infos
|
||||
and len(self._task_infos[self._rcvd_idx]) == 3
|
||||
):
|
||||
info = self._task_infos.pop(self._rcvd_idx)
|
||||
self._structure_infos.append(info[2])
|
||||
return info[1]
|
||||
|
||||
try:
|
||||
# [ avoid hang ]: main process may blocking at _reader.read_next when
|
||||
# KeyboardInterrupt, we do following tradeoff:
|
||||
# 1. get data with timeout, MP_STATUS_CHECK_INTERVAL(5s) as timeout
|
||||
# default, if KeyboardInterrupt blocking, failed workers will be
|
||||
# checked and raise RuntimeError to quit DataLoader in timeout
|
||||
# exception handling.
|
||||
# 2. if get data timeout and check workers all alive, continue to
|
||||
# get data again
|
||||
data = self._data_queue.get(timeout=self._timeout)
|
||||
except Exception as e:
|
||||
# check if thread done event set when waiting data
|
||||
if self._thread_done_event.is_set():
|
||||
continue
|
||||
|
||||
# check failed workers
|
||||
failed_workers = []
|
||||
for i, w in enumerate(self._workers):
|
||||
if self._worker_status[i] and not w.is_alive():
|
||||
failed_workers.append(w)
|
||||
self._shutdown_worker(i)
|
||||
if len(failed_workers) > 0:
|
||||
self._exit_thread_unexpectedly()
|
||||
pids = ', '.join(str(w.pid) for w in failed_workers)
|
||||
logging.warning(
|
||||
f"DataLoader {len(failed_workers)} workers exit unexpectedly, "
|
||||
f"pids: {pids}"
|
||||
)
|
||||
return
|
||||
|
||||
# get(timeout) will call _poll(timeout) and may raise IOError
|
||||
if isinstance(e, (IOError, queue.Empty)):
|
||||
# continue on timeout to keep getting data from queue
|
||||
continue
|
||||
|
||||
self._exit_thread_unexpectedly()
|
||||
logging.error(
|
||||
f"DataLoader reader thread failed({e}) to read data from "
|
||||
"workers' result queue."
|
||||
)
|
||||
raise e
|
||||
else:
|
||||
if self._dataset_kind == _DatasetKind.ITER and isinstance(
|
||||
data, _IterableDatasetStopIteration
|
||||
):
|
||||
# if a worker get StopIteration, we shutdown this worker,
|
||||
# note that this batch indices to trigger StopIteration
|
||||
# is discard, outstanding batch number should be decrease
|
||||
# and another indices should be put for other workers
|
||||
# may still working.
|
||||
if self._persistent_workers:
|
||||
self._worker_status[data.worker_id] = False
|
||||
else:
|
||||
self._shutdown_worker(data.worker_id)
|
||||
self._batches_outstanding -= 1
|
||||
self._try_put_indices()
|
||||
continue
|
||||
|
||||
idx, batch, structure = data
|
||||
|
||||
if (
|
||||
isinstance(idx, _ResumeIteration)
|
||||
and batch is None
|
||||
and structure is None
|
||||
):
|
||||
return idx
|
||||
|
||||
if isinstance(batch, _WorkerException):
|
||||
self._exit_thread_unexpectedly()
|
||||
batch.reraise()
|
||||
|
||||
if idx == self._rcvd_idx:
|
||||
if idx in self._task_infos:
|
||||
del self._task_infos[idx]
|
||||
self._structure_infos.append(structure)
|
||||
return batch
|
||||
else:
|
||||
self._task_infos[idx] += (batch, structure)
|
||||
continue
|
||||
|
||||
def _try_put_indices(self):
|
||||
assert self._batches_outstanding <= self._outstanding_capacity, (
|
||||
"too many indices have been put to queue"
|
||||
)
|
||||
# In multi-process mode for IterableDataset, _try_put_indices will
|
||||
# be called both in main process(for our implement has blocking queue,
|
||||
# and blocking queue read is in main process) and thread, which may
|
||||
# cause error following error
|
||||
# 1. "ValueError: generator already executing" in next(self._sampler_iter)
|
||||
# 2. re-enter in increase _send_idx
|
||||
# add a lock for threading save, for _try_put_indices is only a slight
|
||||
# function which is not in data reading pipeline, this lock almost no
|
||||
# influence on performance
|
||||
with self._thread_lock:
|
||||
try:
|
||||
indices = next(self._sampler_iter)
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
for i in range(self._num_workers):
|
||||
worker_idx = next(self._workers_idx_cycle)
|
||||
if self._worker_status[worker_idx]:
|
||||
break
|
||||
else:
|
||||
return
|
||||
|
||||
self._indices_queues[worker_idx].put((self._send_idx, indices))
|
||||
self._task_infos[self._send_idx] = (worker_idx,)
|
||||
self._batches_outstanding += 1
|
||||
self._send_idx += 1
|
||||
|
||||
def __del__(self):
|
||||
self._try_shutdown_all()
|
||||
|
||||
def _shutdown_on_exit(self):
|
||||
self._try_shutdown_all(1)
|
||||
|
||||
def __next__(self):
|
||||
if in_profiler_mode():
|
||||
trace_event = profiler.RecordEvent(
|
||||
name="_DataLoaderIterMultiProcess",
|
||||
event_type=profiler.TracerEventType.Dataloader,
|
||||
)
|
||||
trace_event.begin()
|
||||
try:
|
||||
benchmark().check_if_need_record(self)
|
||||
benchmark().before_reader()
|
||||
# _batches_outstanding here record the total batch data number
|
||||
# in 'from after _try_put_indices to beforeoutput data', this
|
||||
# value should be _outstanding_capacity if data is not drained,
|
||||
# if _batches_outstanding is less than _places number, there are
|
||||
# no enough data to generate next output, close blocking_queue and
|
||||
# set _thread_done_event here, py_reader will raise StopIteration,
|
||||
# end workers and indices_queues in StopIteration handling
|
||||
if self._batches_outstanding < len(self._places):
|
||||
if self._persistent_workers:
|
||||
raise StopIteration
|
||||
else:
|
||||
self._thread_done_event.set()
|
||||
self._blocking_queue.close()
|
||||
|
||||
if in_dynamic_mode():
|
||||
data = core.eager.read_next_tensor_list(
|
||||
self._reader.read_next_list()[0]
|
||||
)
|
||||
data = _restore_batch(data, self._structure_infos.pop(0))
|
||||
else:
|
||||
if self._return_list:
|
||||
data = self._reader.read_next_list()
|
||||
for i in range(len(data)):
|
||||
data[i] = data[i]._move_to_list()
|
||||
structs = [
|
||||
self._structure_infos.pop(0)
|
||||
for _ in range(len(self._places))
|
||||
]
|
||||
data = [_restore_batch(d, s) for d, s in zip(data, structs)]
|
||||
# static graph organized data on multi-device with list, if
|
||||
# place number is 1, there is only 1 device, extra the data
|
||||
# from list for devices to be compatible with dygraph mode
|
||||
if len(self._places) == 1:
|
||||
data = data[0]
|
||||
else:
|
||||
data = self._reader.read_next()
|
||||
self._on_output_batch()
|
||||
benchmark().after_reader()
|
||||
return data
|
||||
except StopIteration:
|
||||
if not self._persistent_workers:
|
||||
self._reader.shutdown()
|
||||
self._try_shutdown_all()
|
||||
raise
|
||||
finally:
|
||||
if in_profiler_mode():
|
||||
trace_event.end()
|
||||
|
||||
def _on_output_batch(self):
|
||||
for _ in range(len(self._places)):
|
||||
self._batches_outstanding -= 1
|
||||
self._try_put_indices()
|
||||
Executable
+723
@@ -0,0 +1,723 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import bisect
|
||||
import math
|
||||
import warnings
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Generic,
|
||||
TypeVar,
|
||||
)
|
||||
|
||||
from typing_extensions import Never, TypeVarTuple, Unpack, overload
|
||||
|
||||
import paddle
|
||||
from paddle.utils.decorator_utils import variadic_tensor_decorator
|
||||
|
||||
from ... import framework
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import (
|
||||
Callable,
|
||||
Generator,
|
||||
Iterable,
|
||||
Iterator,
|
||||
Sequence,
|
||||
)
|
||||
|
||||
from paddle import Tensor
|
||||
|
||||
_T = TypeVar('_T')
|
||||
_Ts = TypeVarTuple('_Ts')
|
||||
|
||||
|
||||
class Dataset(Generic[_T]):
|
||||
"""
|
||||
An abstract class to encapsulate methods and behaviors of datasets.
|
||||
|
||||
All datasets in map-style(dataset samples can be get by a given key)
|
||||
should be a subclass of `paddle.io.Dataset`. All subclasses should
|
||||
implement following methods:
|
||||
|
||||
:code:`__getitem__`: get sample from dataset with a given index. This
|
||||
method is required by reading dataset sample in :code:`paddle.io.DataLoader`.
|
||||
|
||||
:code:`__len__`: return dataset sample number. This method is required
|
||||
by some implements of :code:`paddle.io.BatchSampler`
|
||||
|
||||
see :code:`paddle.io.DataLoader`.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import Dataset
|
||||
|
||||
>>> # define a random dataset
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([784]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> dataset = RandomDataset(10)
|
||||
>>> for i in range(len(dataset)):
|
||||
... image, label = dataset[i]
|
||||
... # do something
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def __getitem__(self, idx: int) -> _T:
|
||||
raise NotImplementedError(
|
||||
"'{}' not implement in class {}".format(
|
||||
'__getitem__', self.__class__.__name__
|
||||
)
|
||||
)
|
||||
|
||||
def __len__(self) -> int:
|
||||
raise NotImplementedError(
|
||||
"'{}' not implement in class {}".format(
|
||||
'__len__', self.__class__.__name__
|
||||
)
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# A virtual method for type checking only
|
||||
def __iter__(self) -> Iterator[_T]: ...
|
||||
|
||||
|
||||
class IterableDataset(Dataset[_T]):
|
||||
"""
|
||||
An abstract class to encapsulate methods and behaviors of iterable datasets.
|
||||
|
||||
All datasets in iterable-style (can only get sample one by one sequentially, like
|
||||
a Python iterator) should be a subclass of :ref:`api_paddle_io_IterableDataset` . All subclasses should
|
||||
implement following methods:
|
||||
|
||||
:code:`__iter__`: yield sample sequentially. This method is required by reading dataset sample in :ref:`api_paddle_io_DataLoader` .
|
||||
|
||||
.. note::
|
||||
do not implement :code:`__getitem__` and :code:`__len__` in IterableDataset, should not be called either.
|
||||
|
||||
see :ref:`api_paddle_io_DataLoader` .
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: code-example1
|
||||
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import IterableDataset
|
||||
|
||||
>>> # define a random dataset
|
||||
>>> class RandomDataset(IterableDataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __iter__(self):
|
||||
... for i in range(self.num_samples):
|
||||
... image = np.random.random([784]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... yield image, label
|
||||
>>> dataset = RandomDataset(10)
|
||||
>>> for img, label in dataset:
|
||||
... # do something
|
||||
... ...
|
||||
|
||||
When :attr:`num_workers > 0`, each worker has a different copy of the dataset object and
|
||||
will yield whole dataset samples, which means samples in dataset will be repeated in
|
||||
:attr:`num_workers` times. If it is required for each sample to yield only once, there
|
||||
are two methods to configure different copy in each worker process to avoid duplicate data
|
||||
among workers as follows. In both the methods, worker information that can be getted in
|
||||
a worker process by `paddle.io.get_worker_info` will be needed.
|
||||
|
||||
splitting data copy in each worker in :code:`__iter__`
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: code-example2
|
||||
|
||||
>>> import math
|
||||
>>> import paddle
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import IterableDataset, DataLoader, get_worker_info
|
||||
|
||||
>>> class SplitedIterableDataset(IterableDataset): # type: ignore[type-arg]
|
||||
... def __init__(self, start, end):
|
||||
... self.start = start
|
||||
... self.end = end
|
||||
...
|
||||
... def __iter__(self):
|
||||
... worker_info = get_worker_info()
|
||||
... if worker_info is None:
|
||||
... iter_start = self.start
|
||||
... iter_end = self.end
|
||||
... else:
|
||||
... per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
|
||||
... worker_id = worker_info.id
|
||||
... iter_start = self.start + worker_id * per_worker
|
||||
... iter_end = min(iter_start + per_worker, self.end)
|
||||
...
|
||||
... for i in range(iter_start, iter_end):
|
||||
... yield np.array([i])
|
||||
>>> dataset = SplitedIterableDataset(start=2, end=9)
|
||||
>>> dataloader = DataLoader(
|
||||
... dataset,
|
||||
... num_workers=2,
|
||||
... batch_size=1,
|
||||
... drop_last=True,
|
||||
... )
|
||||
>>> for data in dataloader:
|
||||
... print(data) # doctest: +SKIP("The output depends on the environment.")
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[2]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[3]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[4]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[5]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[6]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[7]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[8]])
|
||||
|
||||
splitting data copy in each worker by :code:`worker_init_fn`
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: code-example3
|
||||
|
||||
>>> import math
|
||||
>>> import paddle
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import IterableDataset, DataLoader, get_worker_info
|
||||
|
||||
>>> class RangeIterableDataset(IterableDataset): # type: ignore[type-arg]
|
||||
... def __init__(self, start, end):
|
||||
... self.start = start
|
||||
... self.end = end
|
||||
...
|
||||
... def __iter__(self):
|
||||
... for i in range(self.start, self.end):
|
||||
... yield np.array([i])
|
||||
>>> dataset = RangeIterableDataset(start=2, end=9)
|
||||
|
||||
>>> def worker_init_fn(worker_id):
|
||||
... worker_info = get_worker_info()
|
||||
...
|
||||
... dataset: RangeIterableDataset = worker_info.dataset # type: ignore[assignment]
|
||||
... start = dataset.start
|
||||
... end = dataset.end
|
||||
... num_per_worker = int(math.ceil((end - start) / float(worker_info.num_workers)))
|
||||
...
|
||||
... worker_id = worker_info.id
|
||||
... dataset.start = start + worker_id * num_per_worker
|
||||
... dataset.end = min(dataset.start + num_per_worker, end)
|
||||
>>> dataloader = DataLoader(
|
||||
... dataset,
|
||||
... num_workers=2,
|
||||
... batch_size=1,
|
||||
... drop_last=True,
|
||||
... worker_init_fn=worker_init_fn,
|
||||
... )
|
||||
>>> for data in dataloader:
|
||||
... print(data) # doctest: +SKIP("The output depends on the environment.")
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[2]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[3]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[4]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[5]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[6]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[7]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[8]])
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def __iter__(self) -> Iterator[_T]:
|
||||
raise NotImplementedError(
|
||||
"'{}' not implement in class {}".format(
|
||||
'__iter__', self.__class__.__name__
|
||||
)
|
||||
)
|
||||
|
||||
def __getitem__(self, idx: int) -> Never:
|
||||
raise RuntimeError(
|
||||
"'{}' should not be called for IterableDataset{}".format(
|
||||
'__getitem__', self.__class__.__name__
|
||||
)
|
||||
)
|
||||
|
||||
def __len__(self) -> Never:
|
||||
raise RuntimeError(
|
||||
"'{}' should not be called for IterableDataset{}".format(
|
||||
'__len__', self.__class__.__name__
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TensorDataset(Dataset["Tensor"]):
|
||||
"""
|
||||
Dataset defined by a list of tensors.
|
||||
|
||||
Each tensor should be in shape of [N, ...], while N is the sample number,
|
||||
and each tensor contains a field of sample, :code:`TensorDataset` retrieve
|
||||
each sample by indexing tensors in the 1st dimension.
|
||||
|
||||
Args:
|
||||
tensors(list|tuple): A list/tuple of tensors with same shape in the 1st dimension.
|
||||
|
||||
Returns:
|
||||
Dataset: a Dataset instance wrapping tensors.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> from paddle.io import TensorDataset
|
||||
|
||||
|
||||
>>> input_np = np.random.random([2, 3, 4]).astype('float32')
|
||||
>>> input = paddle.to_tensor(input_np)
|
||||
>>> label_np = np.random.random([2, 1]).astype('int32')
|
||||
>>> label = paddle.to_tensor(label_np)
|
||||
|
||||
>>> dataset = TensorDataset([input, label])
|
||||
|
||||
>>> for i in range(len(dataset)):
|
||||
... input, label = dataset[i]
|
||||
... # do something
|
||||
"""
|
||||
|
||||
tensors: Sequence[Tensor]
|
||||
|
||||
@overload
|
||||
def __init__(self, tensors: Sequence[Tensor]) -> None: ...
|
||||
|
||||
@overload
|
||||
def __init__(self, *tensors: Tensor) -> None: ...
|
||||
|
||||
@variadic_tensor_decorator('tensors', 1)
|
||||
def __init__(self, tensors: Sequence[Tensor]) -> None:
|
||||
if not framework.in_dynamic_mode():
|
||||
raise RuntimeError(
|
||||
"TensorDataset con only be used in imperative mode"
|
||||
)
|
||||
assert all(
|
||||
tensor.shape[0] == tensors[0].shape[0] for tensor in tensors
|
||||
), "tensors not have same shape of the 1st dimension"
|
||||
self.tensors = tensors
|
||||
|
||||
def __getitem__(self, index: int) -> tuple[Tensor, ...]:
|
||||
return tuple(tensor[index] for tensor in self.tensors)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.tensors[0].shape[0]
|
||||
|
||||
|
||||
def to_list(value):
|
||||
if value is None:
|
||||
return value
|
||||
if isinstance(value, (list, tuple)):
|
||||
return list(value)
|
||||
return [value]
|
||||
|
||||
|
||||
class ComposeDataset(Dataset[tuple[Unpack[_Ts]]]):
|
||||
"""
|
||||
A Dataset which composes fields of multiple datasets.
|
||||
|
||||
This dataset is used for composing fields of multiple map-style
|
||||
datasets of same length.
|
||||
|
||||
Args:
|
||||
datasets(list of Dataset): List of datasets to be composed.
|
||||
|
||||
Returns:
|
||||
Dataset: A Dataset which composes fields of multiple datasets.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> from paddle.io import Dataset, ComposeDataset
|
||||
|
||||
>>> # define a random dataset
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([32]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> dataset = ComposeDataset([RandomDataset(10), RandomDataset(10)]) # type: ignore[var-annotated]
|
||||
>>> for i in range(len(dataset)):
|
||||
... image1, label1, image2, label2 = dataset[i]
|
||||
... # do something
|
||||
"""
|
||||
|
||||
datasets: list[Dataset[Any]]
|
||||
|
||||
def __init__(self, datasets: list[Dataset[Any]]) -> None:
|
||||
self.datasets = list(datasets)
|
||||
assert len(self.datasets) > 0, "input datasets should not be empty"
|
||||
for i, dataset in enumerate(self.datasets):
|
||||
assert isinstance(dataset, Dataset), (
|
||||
"each input dataset should be paddle.io.Dataset"
|
||||
)
|
||||
assert not isinstance(dataset, IterableDataset), (
|
||||
"paddle.io.IterableDataset not supported"
|
||||
)
|
||||
if i > 0:
|
||||
assert len(dataset) == len(self.datasets[i - 1]), (
|
||||
"lengths of datasets should be same"
|
||||
)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.datasets[0])
|
||||
|
||||
def __getitem__(self, idx) -> tuple[Unpack[_Ts]]:
|
||||
sample = []
|
||||
for dataset in self.datasets:
|
||||
sample.extend(to_list(dataset[idx]))
|
||||
return tuple(sample)
|
||||
|
||||
|
||||
class ChainDataset(IterableDataset[Any]):
|
||||
"""
|
||||
A Dataset which chains multiple iterable-style datasets.
|
||||
|
||||
This dataset is used for assembling multiple datasets which should
|
||||
be :ref:`api_paddle_io_IterableDataset`.
|
||||
|
||||
Args:
|
||||
datasets(list of IterableDatasets): List of datasets to be chainned.
|
||||
|
||||
Returns:
|
||||
paddle.io.IterableDataset: A Dataset which chains fields of multiple datasets.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> from paddle.io import IterableDataset, ChainDataset
|
||||
|
||||
|
||||
>>> # define a random dataset
|
||||
>>> class RandomDataset(IterableDataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __iter__(self):
|
||||
... for i in range(10):
|
||||
... image = np.random.random([32]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... yield image, label
|
||||
>>> dataset = ChainDataset([RandomDataset(10), RandomDataset(10)])
|
||||
>>> for image, label in iter(dataset):
|
||||
... # do something
|
||||
... ...
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, datasets: list[IterableDataset[Any]]):
|
||||
self.datasets = list(datasets)
|
||||
assert len(self.datasets) > 0, "input datasets should not be empty"
|
||||
for i, dataset in enumerate(self.datasets):
|
||||
assert isinstance(dataset, IterableDataset), (
|
||||
"ChainDataset only support paddle.io.IterableDataset"
|
||||
)
|
||||
|
||||
def __iter__(self) -> Iterator[Any]:
|
||||
for dataset in self.datasets:
|
||||
yield from dataset
|
||||
|
||||
|
||||
class Subset(Dataset[_T]):
|
||||
"""
|
||||
Subset of a dataset at specified indices.
|
||||
|
||||
Args:
|
||||
dataset (Dataset): The whole Dataset.
|
||||
indices (sequence): Indices in the whole set selected for subset.
|
||||
|
||||
Returns:
|
||||
List[Dataset]: A Dataset which is the subset of the original dataset.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
>>> class RangeDataset(paddle.io.Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, start, stop):
|
||||
... self.start = start
|
||||
... self.stop = stop
|
||||
...
|
||||
... def __getitem__(self, index):
|
||||
... return index + self.start
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.stop - self.start
|
||||
|
||||
>>> # Example 1:
|
||||
>>> a = paddle.io.Subset(dataset=RangeDataset(1, 4), indices=[0, 2])
|
||||
>>> print(list(a))
|
||||
[1, 3]
|
||||
|
||||
>>> # Example 2:
|
||||
>>> b = paddle.io.Subset(dataset=RangeDataset(1, 4), indices=[1, 1])
|
||||
>>> print(list(b))
|
||||
[2, 2]
|
||||
"""
|
||||
|
||||
dataset: Dataset[_T]
|
||||
indices: Sequence[int]
|
||||
|
||||
def __init__(self, dataset: Dataset[_T], indices: Sequence[int]) -> None:
|
||||
self.dataset = dataset
|
||||
self.indices = indices
|
||||
|
||||
def __getitem__(self, idx: int) -> _T:
|
||||
return self.dataset[self.indices[idx]]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.indices)
|
||||
|
||||
|
||||
def random_split(
|
||||
dataset: Dataset[_T],
|
||||
lengths: Sequence[int],
|
||||
generator: Any | None = None,
|
||||
) -> list[Subset[_T]]:
|
||||
"""
|
||||
Randomly split a dataset into non-overlapping new datasets of given lengths.
|
||||
Optionally fix the generator for reproducible results, e.g.:
|
||||
|
||||
Args:
|
||||
dataset (Dataset): Dataset to be split
|
||||
lengths (sequence): lengths or fractions of splits to be produced
|
||||
generator (Generator, optional): Generator used for the random permutation. Default is None then the DefaultGenerator is used in manual_seed().
|
||||
|
||||
Returns:
|
||||
Datasets: A list of subset Datasets, which are the non-overlapping subsets of the original Dataset.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
>>> paddle.seed(2023)
|
||||
>>> a_list = paddle.io.random_split(range(10), [3, 7]) # type: ignore[arg-type, var-annotated]
|
||||
>>> print(len(a_list))
|
||||
2
|
||||
|
||||
>>> # output of the first subset
|
||||
>>> for idx, v in enumerate(a_list[0]):
|
||||
... print(idx, v) # doctest: +SKIP("The output depends on the environment.")
|
||||
0 7
|
||||
1 6
|
||||
2 5
|
||||
|
||||
>>> # output of the second subset
|
||||
>>> for idx, v in enumerate(a_list[1]):
|
||||
... print(idx, v) # doctest: +SKIP("The output depends on the environment.")
|
||||
0 1
|
||||
1 9
|
||||
2 4
|
||||
3 2
|
||||
4 0
|
||||
5 3
|
||||
6 8
|
||||
"""
|
||||
if math.isclose(sum(lengths), 1) and sum(lengths) <= 1:
|
||||
subset_lengths = []
|
||||
for i, frac in enumerate(lengths):
|
||||
if frac < 0 or frac > 1:
|
||||
raise ValueError(
|
||||
f"Fraction at index {i} is not between 0 and 1"
|
||||
)
|
||||
n_items_in_split = int(math.floor(len(dataset) * frac))
|
||||
subset_lengths.append(n_items_in_split)
|
||||
remainder = len(dataset) - sum(subset_lengths)
|
||||
|
||||
for i in range(remainder):
|
||||
idx_to_add_at = i % len(subset_lengths)
|
||||
subset_lengths[idx_to_add_at] += 1
|
||||
lengths = subset_lengths
|
||||
for i, length in enumerate(lengths):
|
||||
if length == 0:
|
||||
warnings.warn(
|
||||
f"Length of split at index {i} is 0. "
|
||||
f"This might result in an empty dataset."
|
||||
)
|
||||
|
||||
# Cannot verify that dataset is Sized
|
||||
if sum(lengths) != len(dataset): # type: ignore
|
||||
raise ValueError(
|
||||
"Sum of input lengths does not equal the length of the input dataset!"
|
||||
)
|
||||
# TODO(@Joejiong): support Variable or Tensor type with .tolist class member function.
|
||||
# For example var.item() and var.tolist()
|
||||
indices = paddle.randperm(sum(lengths)).tolist()
|
||||
return [
|
||||
Subset(dataset, indices[offset - length : offset])
|
||||
for offset, length in zip(_accumulate(lengths), lengths)
|
||||
]
|
||||
|
||||
|
||||
def _accumulate(
|
||||
iterable: Iterable[_T], fn: Callable[[_T, _T], _T] = lambda x, y: x + y
|
||||
) -> Generator[_T, None, None]:
|
||||
"""
|
||||
Return running totals
|
||||
|
||||
Args:
|
||||
iterable: any iterable object for example dataset.
|
||||
y (x): one element in the iterable object.
|
||||
fn (x, y): Defaults to lambdax.
|
||||
|
||||
Yields:
|
||||
yields total from beginning iterator to current iterator.
|
||||
|
||||
Example code:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> list(_accumulate([1, 2, 3, 4, 5]))
|
||||
[1, 3, 6, 10, 15]
|
||||
|
||||
>>> import operator
|
||||
>>> list(_accumulate([1, 2, 3, 4, 5], operator.mul))
|
||||
[1, 2, 6, 24, 120]
|
||||
"""
|
||||
|
||||
it = iter(iterable)
|
||||
try:
|
||||
total = next(it)
|
||||
except StopIteration:
|
||||
return
|
||||
yield total
|
||||
for element in it:
|
||||
total = fn(total, element)
|
||||
yield total
|
||||
|
||||
|
||||
class ConcatDataset(Dataset[_T]):
|
||||
"""
|
||||
Dataset as a concatenation of multiple datasets.
|
||||
|
||||
This class is useful to assemble different existing datasets.
|
||||
|
||||
Args:
|
||||
datasets (sequence): List of datasets to be concatenated
|
||||
|
||||
Returns:
|
||||
Dataset: A Dataset which concatenated by multiple datasets.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> from paddle.io import Dataset, ConcatDataset
|
||||
|
||||
>>> # define a random dataset
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([32]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> dataset = ConcatDataset([RandomDataset(10), RandomDataset(10)]) # type: ignore[var-annotated]
|
||||
>>> for i in range(len(dataset)):
|
||||
... image, label = dataset[i]
|
||||
... # do something
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def cumsum(sequence: Sequence[Any]) -> list[int]:
|
||||
r, s = [], 0
|
||||
for e in sequence:
|
||||
l = len(e)
|
||||
r.append(l + s)
|
||||
s += l
|
||||
return r
|
||||
|
||||
def __init__(self, datasets: Iterable[Dataset[Any]]) -> None:
|
||||
self.datasets = list(datasets)
|
||||
assert len(self.datasets) > 0, (
|
||||
'datasets should not be an empty iterable'
|
||||
)
|
||||
for d in self.datasets:
|
||||
assert not isinstance(d, IterableDataset), (
|
||||
"ConcatDataset does not support IterableDataset"
|
||||
)
|
||||
self.cumulative_sizes = self.cumsum(self.datasets)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.cumulative_sizes[-1]
|
||||
|
||||
def __getitem__(self, idx: int) -> _T:
|
||||
if idx < 0:
|
||||
if -idx > len(self):
|
||||
raise ValueError(
|
||||
"absolute value of index should not exceed dataset length"
|
||||
)
|
||||
idx = len(self) + idx
|
||||
dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
|
||||
if dataset_idx == 0:
|
||||
sample_idx = idx
|
||||
else:
|
||||
sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
|
||||
return self.datasets[dataset_idx][sample_idx]
|
||||
@@ -0,0 +1,86 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
class _DatasetFetcher:
|
||||
def __init__(self, dataset, auto_collate_batch, collate_fn, drop_last):
|
||||
self.dataset = dataset
|
||||
self.auto_collate_batch = auto_collate_batch
|
||||
self.collate_fn = collate_fn
|
||||
self.drop_last = drop_last
|
||||
|
||||
# NOTE: fetch function here perform the whole pipeline of dataset
|
||||
# reading and data transforms of a batch in each calling, this
|
||||
# may take a long time inside, if DataLoader is exit outside,
|
||||
# fetch need to perceive exit situation, so we pass done_event
|
||||
# here for fetch to check exit status
|
||||
# NOTE: if DataLoader exit by `break`, performing GPU tensor operations,
|
||||
# e.g. to_tensor may cause SIGSEGV in thread, so we pass the
|
||||
# done_event argument to check DataLoader exit status between
|
||||
# each sample processing in the batch
|
||||
def fetch(self, batch_indices, done_event=None):
|
||||
raise NotImplementedError(
|
||||
f"'fetch' not implement for class {self.__class__.__name__}"
|
||||
)
|
||||
|
||||
|
||||
class _IterableDatasetFetcher(_DatasetFetcher):
|
||||
def __init__(self, dataset, auto_collate_batch, collate_fn, drop_last):
|
||||
super().__init__(dataset, auto_collate_batch, collate_fn, drop_last)
|
||||
self.dataset_iter = iter(dataset)
|
||||
|
||||
def fetch(self, batch_indices, done_event=None):
|
||||
if self.auto_collate_batch:
|
||||
data = []
|
||||
for _ in batch_indices:
|
||||
if done_event is None or not done_event.is_set():
|
||||
try:
|
||||
data.append(next(self.dataset_iter))
|
||||
except StopIteration:
|
||||
break
|
||||
else:
|
||||
return None
|
||||
|
||||
if len(data) == 0 or (
|
||||
self.drop_last and len(data) < len(batch_indices)
|
||||
):
|
||||
raise StopIteration
|
||||
|
||||
else:
|
||||
data = next(self.dataset_iter)
|
||||
|
||||
if self.collate_fn:
|
||||
data = self.collate_fn(data)
|
||||
return data
|
||||
|
||||
|
||||
class _MapDatasetFetcher(_DatasetFetcher):
|
||||
def __init__(self, dataset, auto_collate_batch, collate_fn, drop_last):
|
||||
super().__init__(dataset, auto_collate_batch, collate_fn, drop_last)
|
||||
|
||||
def fetch(self, batch_indices, done_event=None):
|
||||
if self.auto_collate_batch:
|
||||
data = []
|
||||
for idx in batch_indices:
|
||||
if done_event is None or not done_event.is_set():
|
||||
data.append(self.dataset[idx])
|
||||
else:
|
||||
return None
|
||||
|
||||
else:
|
||||
data = self.dataset[batch_indices]
|
||||
|
||||
if self.collate_fn:
|
||||
data = self.collate_fn(data)
|
||||
return data
|
||||
@@ -0,0 +1,152 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import numbers
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
|
||||
FIELD_PREFIX = "_paddle_field_"
|
||||
|
||||
|
||||
def _flatten_batch(batch):
|
||||
"""
|
||||
For lod_blocking_queue only receive tensor array, flatten batch
|
||||
data, extract numpy.array data out as a list of numpy.array to
|
||||
send to lod_blocking_queue, and save the batch data structure
|
||||
such as fields in other types (str, int, etc) or key-value map
|
||||
of dictionaries
|
||||
"""
|
||||
|
||||
def _flatten(batch, flat_batch, structure, field_idx):
|
||||
if isinstance(batch, Sequence):
|
||||
for field in batch:
|
||||
if isinstance(
|
||||
field,
|
||||
(np.ndarray, paddle.Tensor, paddle.base.core.eager.Tensor),
|
||||
):
|
||||
structure.append(f'{FIELD_PREFIX}{field_idx}')
|
||||
flat_batch.append(field)
|
||||
field_idx += 1
|
||||
elif isinstance(field, (str, bytes, numbers.Number)):
|
||||
structure.append(field)
|
||||
elif isinstance(field, Sequence):
|
||||
field_struct, field_idx = _flatten(
|
||||
field, flat_batch, [], field_idx
|
||||
)
|
||||
structure.append(field_struct)
|
||||
elif isinstance(field, Mapping):
|
||||
field_struct, field_idx = _flatten(
|
||||
field, flat_batch, {}, field_idx
|
||||
)
|
||||
structure.append(field_struct)
|
||||
else:
|
||||
structure.append(field)
|
||||
elif isinstance(batch, Mapping):
|
||||
for k, field in batch.items():
|
||||
if isinstance(
|
||||
field,
|
||||
(np.ndarray, paddle.Tensor, paddle.base.core.eager.Tensor),
|
||||
):
|
||||
structure[k] = f'{FIELD_PREFIX}{field_idx}'
|
||||
flat_batch.append(field)
|
||||
field_idx += 1
|
||||
elif isinstance(field, (str, bytes, numbers.Number)):
|
||||
structure[k] = field
|
||||
elif isinstance(field, Sequence):
|
||||
field_struct, field_idx = _flatten(
|
||||
field, flat_batch, [], field_idx
|
||||
)
|
||||
structure[k] = field_struct
|
||||
elif isinstance(field, Mapping):
|
||||
field_struct, field_idx = _flatten(
|
||||
field, flat_batch, {}, field_idx
|
||||
)
|
||||
structure[k] = field_struct
|
||||
else:
|
||||
structure[k] = field
|
||||
else:
|
||||
raise TypeError(f"wrong flat data type: {type(batch)}")
|
||||
|
||||
return structure, field_idx
|
||||
|
||||
# sample only contains single fields
|
||||
if not isinstance(batch, Sequence):
|
||||
flat_batch = []
|
||||
structure, _ = _flatten([batch], flat_batch, [], 0)
|
||||
return flat_batch, structure[0]
|
||||
flat_batch = []
|
||||
structure, _ = _flatten(batch, flat_batch, [], 0)
|
||||
return flat_batch, structure
|
||||
|
||||
|
||||
def _restore_batch(flat_batch, structure):
|
||||
"""
|
||||
After reading list of Tensor data from lod_blocking_queue outputs,
|
||||
use this function to restore the batch data structure, replace
|
||||
:attr:`_paddle_field_x` with data from flat_batch
|
||||
"""
|
||||
|
||||
def _restore(structure, field_idx):
|
||||
if isinstance(structure, Sequence):
|
||||
for i, field in enumerate(structure):
|
||||
if isinstance(field, str) and field.startswith(FIELD_PREFIX):
|
||||
cur_field_idx = int(field.replace(FIELD_PREFIX, ''))
|
||||
field_idx = max(field_idx, cur_field_idx)
|
||||
assert flat_batch[cur_field_idx] is not None, (
|
||||
"flat_batch[{}] parsed repeatedly"
|
||||
)
|
||||
structure[i] = flat_batch[cur_field_idx]
|
||||
flat_batch[cur_field_idx] = None
|
||||
elif isinstance(field, (str, bytes, numbers.Number)):
|
||||
continue
|
||||
elif isinstance(field, (Sequence, Mapping)):
|
||||
field_idx = _restore(structure[i], field_idx)
|
||||
elif isinstance(structure, Mapping):
|
||||
for k, field in structure.items():
|
||||
if isinstance(field, str) and field.startswith(FIELD_PREFIX):
|
||||
cur_field_idx = int(field.replace(FIELD_PREFIX, ''))
|
||||
field_idx = max(field_idx, cur_field_idx)
|
||||
assert flat_batch[cur_field_idx] is not None, (
|
||||
"flat_batch[{}] parsed repeatedly"
|
||||
)
|
||||
structure[k] = flat_batch[cur_field_idx]
|
||||
flat_batch[cur_field_idx] = None
|
||||
elif isinstance(field, (str, bytes, numbers.Number)):
|
||||
continue
|
||||
elif isinstance(field, (Sequence, Mapping)):
|
||||
field_idx = _restore(structure[k], field_idx)
|
||||
else:
|
||||
raise TypeError(f"wrong flat data type: {type(structure)}")
|
||||
|
||||
return field_idx
|
||||
|
||||
assert isinstance(flat_batch, Sequence), "flat_batch is not a list or tuple"
|
||||
|
||||
# no np.array in dataset, no output tensor from blocking queue
|
||||
# simply return structure
|
||||
if len(flat_batch) == 0:
|
||||
return structure
|
||||
|
||||
# sample only contains single fields
|
||||
if isinstance(structure, (str, bytes)):
|
||||
assert structure == f'{FIELD_PREFIX}{0}', (
|
||||
f"invalid structure: {structure}"
|
||||
)
|
||||
return flat_batch[0]
|
||||
field_idx = _restore(structure, 0)
|
||||
assert field_idx + 1 == len(flat_batch), "Tensor parse incomplete"
|
||||
return structure
|
||||
@@ -0,0 +1,436 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from collections.abc import Iterator
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Generic,
|
||||
TypeVar,
|
||||
)
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ...framework import core
|
||||
from ...tensor import randperm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Generator, Sequence, Sized
|
||||
|
||||
import numpy.typing as npt
|
||||
|
||||
from paddle import Tensor
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
class Sampler(Generic[_T]):
|
||||
"""
|
||||
An abstract class to encapsulate methods and behaviors of samplers.
|
||||
|
||||
All sampler used by :code:`paddle.io.BatchSampler` should be a subclass
|
||||
of :code:`paddle.io.Sampler`, BatchSampler subclasses should
|
||||
implement following methods:
|
||||
|
||||
:code:`__iter__`: return sample index iterably, which iterate over indices
|
||||
of dataset elements
|
||||
|
||||
:code:`__len__`: the number of sample in :attr:`data_source`
|
||||
|
||||
|
||||
Args:
|
||||
data_source(Dataset, optional): this could be an instance of
|
||||
:code:`paddle.io.Dataset` other Python object which
|
||||
implemented :code:`__len__` for Sampler to get indices
|
||||
as the range of :attr:`dataset` length. Default None.
|
||||
|
||||
Returns:
|
||||
Sampler: an iterable object for sample indices iterating
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import Dataset, Sampler
|
||||
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([784]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> class MySampler(Sampler): # type: ignore[type-arg]
|
||||
... def __init__(self, data_source):
|
||||
... self.data_source = data_source
|
||||
...
|
||||
... def __iter__(self):
|
||||
... return iter(range(len(self.data_source))) # type: ignore[arg-type]
|
||||
...
|
||||
... def __len__(self):
|
||||
... return len(self.data_source) # type: ignore[arg-type]
|
||||
>>> sampler = MySampler(data_source=RandomDataset(100))
|
||||
|
||||
>>> for index in sampler:
|
||||
... print(index)
|
||||
0
|
||||
1
|
||||
2
|
||||
...
|
||||
99
|
||||
|
||||
see `paddle.io.BatchSampler`
|
||||
see `paddle.io.DataLoader`
|
||||
|
||||
"""
|
||||
|
||||
data_source: Sized | None
|
||||
|
||||
def __init__(self, data_source: Sized | None = None) -> None:
|
||||
self.data_source = data_source
|
||||
|
||||
def __iter__(self) -> Iterator[_T]:
|
||||
raise NotImplementedError
|
||||
|
||||
# Not define __len__ method in this base class here for __len__
|
||||
# is not needed in same sense, e.g. paddle.io.IterableDataset
|
||||
if TYPE_CHECKING:
|
||||
|
||||
def __len__(self) -> int: ...
|
||||
|
||||
|
||||
class SequenceSampler(Sampler[int]):
|
||||
"""
|
||||
Iterate samples sequentially, yield :code:`0, 1, 2, ..., len(data_source) -1`
|
||||
generally,
|
||||
|
||||
Args:
|
||||
data_source(Dataset): dataset to sample, this could be an
|
||||
instance of :code:`paddle.io.Dataset` other Python
|
||||
object which implemented :code:`__len__`.
|
||||
|
||||
Returns:
|
||||
Sampler: a Sampler yield sample index sequentially
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import Dataset, SequenceSampler
|
||||
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([784]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> sampler = SequenceSampler(data_source=RandomDataset(100))
|
||||
|
||||
>>> for index in sampler:
|
||||
... print(index)
|
||||
0
|
||||
1
|
||||
2
|
||||
...
|
||||
99
|
||||
|
||||
see `paddle.io.Sampler`
|
||||
"""
|
||||
|
||||
data_source: Sized
|
||||
|
||||
def __init__(self, data_source: Sized) -> None:
|
||||
self.data_source = data_source
|
||||
|
||||
def __iter__(self) -> Iterator[int]:
|
||||
return iter(range(len(self.data_source)))
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.data_source)
|
||||
|
||||
|
||||
class RandomSampler(Sampler[int]):
|
||||
"""
|
||||
Iterate samples randomly, yield shuffled indices, if :attr:`replacement=False`,
|
||||
yield shuffled indices of the whole data source, if :attr:`replacement=True`,
|
||||
:attr:`num_samples` can set to specify the sample number to draw.
|
||||
|
||||
Args:
|
||||
data_source(Dataset): dataset to sample, this could be an
|
||||
instance of :ref:`api_paddle_io_Dataset` or :ref:`api_paddle_io_IterableDataset` or other Python
|
||||
object which implemented :code:`__len__` to get indices as the range of :code:`dataset` length. Default None.
|
||||
replacement(bool, optional): If False, sample the whole dataset, If True,
|
||||
set :attr:`num_samples` for how many samples to draw. Default False.
|
||||
num_samples(int, optional): set sample number to draw. Default None, which is set to the length of `data_source`.
|
||||
generator(Generator, optional): specify a generator to sample the :code:`data_source`. Default None, disabled.
|
||||
|
||||
Returns:
|
||||
RandomSampler: a Sampler yield sample index randomly.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import Dataset, RandomSampler
|
||||
|
||||
>>> np.random.seed(2023)
|
||||
>>> class RandomDataset(Dataset): # type: ignore[type-arg]
|
||||
... def __init__(self, num_samples):
|
||||
... self.num_samples = num_samples
|
||||
...
|
||||
... def __getitem__(self, idx):
|
||||
... image = np.random.random([784]).astype('float32')
|
||||
... label = np.random.randint(0, 9, (1,)).astype('int64')
|
||||
... return image, label
|
||||
...
|
||||
... def __len__(self):
|
||||
... return self.num_samples
|
||||
>>> sampler = RandomSampler(data_source=RandomDataset(100))
|
||||
|
||||
>>> for index in sampler:
|
||||
... print(index)
|
||||
56
|
||||
12
|
||||
68
|
||||
...
|
||||
87
|
||||
"""
|
||||
|
||||
data_source: Sized
|
||||
replacement: bool
|
||||
generator: Generator[int, None, None] | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_source: Sized,
|
||||
replacement: bool = False,
|
||||
num_samples: int | None = None,
|
||||
generator: Generator[int, None, None] | None = None,
|
||||
) -> None:
|
||||
self.data_source = data_source
|
||||
self.replacement = replacement
|
||||
self._num_samples = num_samples
|
||||
if isinstance(generator, Iterator):
|
||||
self.generator = generator
|
||||
else:
|
||||
warnings.warn(
|
||||
"the specified generator is not iterable and will be ignored"
|
||||
)
|
||||
self.generator = None
|
||||
|
||||
if not isinstance(self.replacement, bool):
|
||||
raise TypeError(
|
||||
"expect boolean value for replacement, but got "
|
||||
f"replacement={self.replacement}"
|
||||
)
|
||||
|
||||
if not self.replacement and self.num_samples > len(self.data_source):
|
||||
raise ValueError(
|
||||
"num_samples should be smaller than or equal to length of data_source when replacement is False, "
|
||||
f"but got num_samples: {self.num_samples} > data_source: {len(self.data_source)}"
|
||||
)
|
||||
|
||||
if not isinstance(self.num_samples, int) or self.num_samples <= 0:
|
||||
raise ValueError(
|
||||
"num_samples should be a positive integer, "
|
||||
f"but got num_samples={self.num_samples}"
|
||||
)
|
||||
|
||||
@property
|
||||
def num_samples(self) -> int:
|
||||
if self._num_samples is None:
|
||||
return len(self.data_source)
|
||||
return self._num_samples
|
||||
|
||||
def __iter__(self) -> Iterator[int]:
|
||||
n = len(self.data_source)
|
||||
if self.generator:
|
||||
for i in range(self.num_samples):
|
||||
try:
|
||||
index = next(self.generator)
|
||||
except StopIteration:
|
||||
return
|
||||
yield index
|
||||
else:
|
||||
if self.replacement:
|
||||
for index in np.random.choice(
|
||||
np.arange(n), self.num_samples, replace=True
|
||||
).tolist():
|
||||
yield index
|
||||
else:
|
||||
for index in np.random.choice(
|
||||
np.arange(n), self.num_samples, replace=False
|
||||
).tolist():
|
||||
yield index
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_samples
|
||||
|
||||
|
||||
def _weighted_sample(weights, num_samples, replacement=True):
|
||||
if isinstance(weights, core.DenseTensor):
|
||||
weights = weights.numpy()
|
||||
if isinstance(weights, (list, tuple)):
|
||||
weights = np.array(weights)
|
||||
assert isinstance(weights, np.ndarray), (
|
||||
"weights should be paddle.Tensor, numpy.ndarray, list or tuple"
|
||||
)
|
||||
assert len(weights.shape) <= 2, "weights should be a 1-D or 2-D array"
|
||||
weights = weights.reshape((-1, weights.shape[-1]))
|
||||
assert np.all(weights >= 0.0), "weights should be positive value"
|
||||
assert not np.any(weights == np.inf), "weights should not be INF"
|
||||
assert not np.any(weights == np.nan), "weights should not be NaN"
|
||||
|
||||
non_zeros = np.sum(weights > 0.0, axis=1)
|
||||
assert np.all(non_zeros > 0), "weights should have positive values"
|
||||
if not replacement:
|
||||
assert np.all(non_zeros >= num_samples), (
|
||||
"weights positive value number should not "
|
||||
"less than num_samples when replacement=False"
|
||||
)
|
||||
|
||||
weights = weights / weights.sum(axis=1)
|
||||
rets = []
|
||||
for i in range(weights.shape[0]):
|
||||
ret = np.random.choice(
|
||||
weights.shape[1], num_samples, replacement, weights[i]
|
||||
)
|
||||
rets.append(ret)
|
||||
return np.array(rets)
|
||||
|
||||
|
||||
class WeightedRandomSampler(Sampler[int]):
|
||||
"""
|
||||
Random sample with given weights (probabilities), sample index will be in range
|
||||
[0, len(weights) - 1], if :attr:`replacement` is True, index can be sampled
|
||||
multiple times.
|
||||
|
||||
Args:
|
||||
weights(numpy.ndarray|paddle.Tensor|list|tuple): sequence of weights,
|
||||
should be numpy array, paddle.Tensor, list or tuple
|
||||
num_samples(int): set sample number to draw from sampler.
|
||||
replacement(bool): Whether to draw sample with replacements, default True
|
||||
|
||||
Returns:
|
||||
Sampler: a Sampler yield sample index randomly by given weights
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import WeightedRandomSampler
|
||||
|
||||
>>> np.random.seed(2023)
|
||||
>>> sampler = WeightedRandomSampler(
|
||||
... weights=[0.1, 0.3, 0.5, 0.7, 0.2],
|
||||
... num_samples=5,
|
||||
... replacement=True,
|
||||
... )
|
||||
>>> for index in sampler:
|
||||
... print(index)
|
||||
2
|
||||
4
|
||||
3
|
||||
1
|
||||
1
|
||||
"""
|
||||
|
||||
weights: npt.NDArray[Any] | Tensor | Sequence[float]
|
||||
num_samples: int
|
||||
replacement: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weights: npt.NDArray[Any] | Tensor | Sequence[float],
|
||||
num_samples: int,
|
||||
replacement: bool = True,
|
||||
) -> None:
|
||||
if not isinstance(num_samples, int) or num_samples <= 0:
|
||||
raise ValueError("num_samples should be a positive integer")
|
||||
if not isinstance(replacement, bool):
|
||||
raise ValueError("replacement should be a boolean value")
|
||||
self.weights = weights
|
||||
self.num_samples = num_samples
|
||||
self.replacement = replacement
|
||||
|
||||
def __iter__(self) -> Iterator[int]:
|
||||
idxs = _weighted_sample(
|
||||
self.weights, self.num_samples, self.replacement
|
||||
)
|
||||
return iter(idxs.reshape(-1).tolist())
|
||||
|
||||
def __len__(self) -> int:
|
||||
mul = np.prod(self.weights.shape) // self.weights.shape[-1]
|
||||
return self.num_samples * mul
|
||||
|
||||
|
||||
class SubsetRandomSampler(Sampler[int]):
|
||||
r"""
|
||||
Randomly sample elements from a given list of indices, without replacement.
|
||||
|
||||
Args:
|
||||
indices (sequence): a sequence of indices
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.io import SubsetRandomSampler
|
||||
|
||||
>>> paddle.seed(2023)
|
||||
>>> sampler = SubsetRandomSampler(indices=[1, 3, 5, 7, 9])
|
||||
|
||||
>>> for index in sampler:
|
||||
... print(index)
|
||||
9
|
||||
3
|
||||
7
|
||||
5
|
||||
1
|
||||
|
||||
"""
|
||||
|
||||
indices: Sequence[int]
|
||||
|
||||
def __init__(self, indices: Sequence[int]) -> None:
|
||||
if len(indices) == 0:
|
||||
raise ValueError(
|
||||
"The length of `indices` in SubsetRandomSampler should be greater than 0."
|
||||
)
|
||||
self.indices = indices
|
||||
|
||||
def __iter__(self) -> Iterator[int]:
|
||||
for i in randperm(len(self.indices)):
|
||||
yield self.indices[i]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.indices)
|
||||
@@ -0,0 +1,419 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import queue
|
||||
import sys
|
||||
import traceback
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
|
||||
from ...framework import core
|
||||
from ..multiprocess_utils import (
|
||||
MP_STATUS_CHECK_INTERVAL,
|
||||
CleanupFuncRegistrar,
|
||||
_cleanup_mmap,
|
||||
)
|
||||
from .fetcher import _IterableDatasetFetcher, _MapDatasetFetcher
|
||||
from .flat import _flatten_batch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle.io import Dataset
|
||||
|
||||
|
||||
class _IterableDatasetStopIteration:
|
||||
def __init__(self, worker_id):
|
||||
self.worker_id = worker_id
|
||||
|
||||
|
||||
class _ResumeIteration:
|
||||
pass
|
||||
|
||||
|
||||
class _DatasetKind:
|
||||
MAP = 0
|
||||
ITER = 1
|
||||
|
||||
@staticmethod
|
||||
def create_fetcher(
|
||||
kind, dataset, auto_collate_batch, collate_fn, drop_last
|
||||
):
|
||||
if kind == _DatasetKind.MAP:
|
||||
return _MapDatasetFetcher(
|
||||
dataset, auto_collate_batch, collate_fn, drop_last
|
||||
)
|
||||
elif kind == _DatasetKind.ITER:
|
||||
return _IterableDatasetFetcher(
|
||||
dataset, auto_collate_batch, collate_fn, drop_last
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"unknown Dataset kind {kind}")
|
||||
|
||||
|
||||
class ParentWatchDog:
|
||||
def __init__(self):
|
||||
self._parent_pid = os.getppid()
|
||||
self._parent_alive = True
|
||||
|
||||
def is_alive(self):
|
||||
if self._parent_alive:
|
||||
self._parent_alive = os.getppid() == self._parent_pid
|
||||
return self._parent_alive
|
||||
|
||||
|
||||
# worker information for each workers, used for splitting data copy
|
||||
# for IteratorDataset in worker processes.
|
||||
_worker_info = None
|
||||
|
||||
|
||||
def get_worker_info() -> WorkerInfo:
|
||||
"""
|
||||
Get DataLoader worker process information function, this function is
|
||||
used to split data copy in worker process for IterableDataset
|
||||
(see :code:`paddle.io.IterableDataset`), worker information contains
|
||||
following fields:
|
||||
|
||||
:attr:`num_workers`: total worker process number, see `paddle.io.DataLoader`
|
||||
|
||||
:attr:`id`: the worker process id, count from 0 to :attr:`num_workers - 1`
|
||||
|
||||
:attr:`dataset`: the dataset object in this worker process
|
||||
|
||||
Returns:
|
||||
WorkerInfo: an instance of WorkerInfo which contains fields above.
|
||||
|
||||
Notes:
|
||||
For more usage and examples, please see :code:`paddle.io.IterableDataset`
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import math
|
||||
>>> import paddle
|
||||
>>> import numpy as np
|
||||
>>> from paddle.io import IterableDataset, DataLoader, get_worker_info
|
||||
|
||||
>>> class SplitedIterableDataset(IterableDataset): # type: ignore[type-arg]
|
||||
... def __init__(self, start, end):
|
||||
... self.start = start
|
||||
... self.end = end
|
||||
...
|
||||
... def __iter__(self):
|
||||
... worker_info = get_worker_info()
|
||||
... if worker_info is None:
|
||||
... iter_start = self.start
|
||||
... iter_end = self.end
|
||||
... else:
|
||||
... per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
|
||||
... worker_id = worker_info.id
|
||||
... iter_start = self.start + worker_id * per_worker
|
||||
... iter_end = min(iter_start + per_worker, self.end)
|
||||
...
|
||||
... for i in range(iter_start, iter_end):
|
||||
... yield np.array([i])
|
||||
>>> place = paddle.CPUPlace()
|
||||
>>> dataset = SplitedIterableDataset(start=2, end=9)
|
||||
>>> dataloader = DataLoader(
|
||||
... dataset,
|
||||
... places=place,
|
||||
... num_workers=2,
|
||||
... batch_size=1,
|
||||
... drop_last=True,
|
||||
... )
|
||||
>>> for data in dataloader:
|
||||
... print(data) # doctest: +SKIP("The output depends on the environment.")
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[2]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[6]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[3]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[7]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[4]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[8]])
|
||||
Tensor(shape=[1, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
|
||||
[[5]])
|
||||
|
||||
"""
|
||||
return _worker_info
|
||||
|
||||
|
||||
class WorkerInfo:
|
||||
num_workers: int
|
||||
id: int
|
||||
dataset: Dataset[Any]
|
||||
seed: int
|
||||
|
||||
__initialized = False
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
for k, v in kwargs.items():
|
||||
setattr(self, k, v)
|
||||
self.__initialized = True
|
||||
|
||||
def __setattr__(self, key, val):
|
||||
if self.__initialized:
|
||||
raise RuntimeError(
|
||||
f"Cannot assign attributes to {self.__class__.__name__} objects"
|
||||
)
|
||||
return super().__setattr__(key, val)
|
||||
|
||||
|
||||
class _WorkerException:
|
||||
def __init__(self, worker_id, exc_info=None):
|
||||
self.worker_id = worker_id
|
||||
exc_info = exc_info or sys.exc_info()
|
||||
self.exc_type = exc_info[0]
|
||||
self.exc_msg = "".join(traceback.format_exception(*exc_info))
|
||||
|
||||
def reraise(self):
|
||||
msg = f"DataLoader worker({self.worker_id}) caught {self.exc_type.__name__} with message:\n{self.exc_msg}"
|
||||
if getattr(self.exc_type, "message", None):
|
||||
raise self.exc_type(message=msg)
|
||||
raise self.exc_type(msg)
|
||||
|
||||
|
||||
# The function `_generate_states` is adapted from `numpy.random.SeedSequence`
|
||||
# from https://github.com/numpy/numpy/blob/main/numpy/random/bit_generator.pyx
|
||||
# Here is the copyright:
|
||||
|
||||
# SeedSequence is derived from Melissa E. O'Neill's C++11 `std::seed_seq`
|
||||
# implementation, as it has a lot of nice properties that we want.
|
||||
# https://gist.github.com/imneme/540829265469e673d045
|
||||
# http://www.pcg-random.org/posts/developing-a-seed_seq-alternative.html
|
||||
|
||||
# The MIT License (MIT)
|
||||
|
||||
# Copyright (c) 2015 Melissa E. O'Neill
|
||||
# Copyright (c) 2019 NumPy Developers
|
||||
#
|
||||
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
# of this software and associated documentation files (the "Software"), to deal
|
||||
# in the Software without restriction, including without limitation the rights
|
||||
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
# copies of the Software, and to permit persons to whom the Software is
|
||||
# furnished to do so, subject to the following conditions:
|
||||
#
|
||||
# The above copyright notice and this permission notice shall be included in
|
||||
# all copies or substantial portions of the Software.
|
||||
#
|
||||
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
# SOFTWARE.
|
||||
|
||||
INIT_A = 0x43B0D7E5
|
||||
MULT_A = 0x931E8875
|
||||
INIT_B = 0x8B51F9DD
|
||||
MULT_B = 0x58F38DED
|
||||
MIX_MULT_L = 0xCA01F9DD
|
||||
MIX_MULT_R = 0x4973F715
|
||||
XSHIFT = np.dtype(np.uint32).itemsize * 8 // 2
|
||||
MASK32 = 0xFFFFFFFF
|
||||
|
||||
|
||||
def _generate_states(base_seed=0, worker_id=0):
|
||||
# init hash constant
|
||||
hash_const_A = INIT_A
|
||||
hash_const_B = INIT_B
|
||||
|
||||
def hash(value):
|
||||
nonlocal hash_const_A
|
||||
value = (value ^ hash_const_A) & MASK32
|
||||
hash_const_A = (hash_const_A * MULT_A) & MASK32
|
||||
value = (value * hash_const_A) & MASK32
|
||||
value = (value ^ (value >> XSHIFT)) & MASK32
|
||||
return value
|
||||
|
||||
def mix(x, y):
|
||||
result_x = (MIX_MULT_L * x) & MASK32
|
||||
result_y = (MIX_MULT_R * y) & MASK32
|
||||
result = (result_x - result_y) & MASK32
|
||||
result = (result ^ (result >> XSHIFT)) & MASK32
|
||||
return result
|
||||
|
||||
# init entropies with based_seed and worker_id and calculate pool
|
||||
entropies = [worker_id, base_seed & MASK32, base_seed >> 32, 0]
|
||||
pool = [hash(entropy) for entropy in entropies]
|
||||
|
||||
# mix all bits together
|
||||
for i in range(len(pool)):
|
||||
for j in range(len(pool)):
|
||||
if i != j:
|
||||
pool[j] = mix(pool[j], hash(pool[i]))
|
||||
|
||||
states = []
|
||||
for p in pool:
|
||||
state = (p ^ hash_const_B) & MASK32
|
||||
hash_const_B = (hash_const_B * MULT_B) & MASK32
|
||||
state = (state * hash_const_B) & MASK32
|
||||
state = (state ^ (state >> XSHIFT)) & MASK32
|
||||
states.append(state)
|
||||
|
||||
return states
|
||||
|
||||
|
||||
def _worker_loop(
|
||||
dataset,
|
||||
dataset_kind,
|
||||
indices_queue,
|
||||
out_queue,
|
||||
done_event,
|
||||
auto_collate_batch,
|
||||
collate_fn,
|
||||
drop_last,
|
||||
init_fn,
|
||||
worker_id,
|
||||
num_workers,
|
||||
use_shared_memory,
|
||||
base_seed,
|
||||
shm_cache_size=0,
|
||||
):
|
||||
try:
|
||||
# NOTE: [ mmap files clear ] When the child process exits unexpectedly,
|
||||
# some shared memory objects may have been applied for but have not yet
|
||||
# been put into the inter-process Queue. This part of the object needs
|
||||
# to be cleaned up when the process ends.
|
||||
CleanupFuncRegistrar.register(_cleanup_mmap)
|
||||
|
||||
# set signal handler
|
||||
core._set_process_signal_handler()
|
||||
|
||||
core._set_max_memory_map_allocation_pool_size(shm_cache_size)
|
||||
|
||||
# set different numpy seed for each worker
|
||||
try:
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
except ImportError:
|
||||
pass
|
||||
else:
|
||||
seed = base_seed + worker_id
|
||||
random.seed(seed)
|
||||
paddle.seed(seed)
|
||||
np.random.seed(_generate_states(base_seed, worker_id))
|
||||
|
||||
global _worker_info
|
||||
_worker_info = WorkerInfo(
|
||||
id=worker_id,
|
||||
num_workers=num_workers,
|
||||
dataset=dataset,
|
||||
seed=base_seed,
|
||||
)
|
||||
|
||||
init_exception = None
|
||||
try:
|
||||
if init_fn is not None:
|
||||
init_fn(worker_id)
|
||||
fetcher = _DatasetKind.create_fetcher(
|
||||
dataset_kind, dataset, auto_collate_batch, collate_fn, drop_last
|
||||
)
|
||||
except:
|
||||
init_exception = _WorkerException(worker_id)
|
||||
|
||||
iterator_drained = False
|
||||
parent_watch_dog = ParentWatchDog()
|
||||
|
||||
while parent_watch_dog.is_alive():
|
||||
try:
|
||||
data = indices_queue.get(MP_STATUS_CHECK_INTERVAL)
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
if isinstance(data, _ResumeIteration):
|
||||
out_queue.put((data, None, None))
|
||||
iterator_drained = False
|
||||
fetcher = _DatasetKind.create_fetcher(
|
||||
dataset_kind, dataset, auto_collate_batch, collate_fn, True
|
||||
)
|
||||
continue
|
||||
|
||||
# None as poison piil, so worker event should be set
|
||||
if data is None:
|
||||
assert done_event.is_set() or iterator_drained, (
|
||||
"get None when worker done_event set"
|
||||
)
|
||||
break
|
||||
# If worker done event is set but get still get data in
|
||||
# indices_queue, remaining data should be get and skipped.
|
||||
if done_event.is_set() or iterator_drained:
|
||||
continue
|
||||
|
||||
idx, indices = data
|
||||
try:
|
||||
if init_exception is not None:
|
||||
batch = init_exception
|
||||
init_exception = None
|
||||
else:
|
||||
# NOTE: GPU tensor operation is not supported in sub-process
|
||||
# but default device is GPU in paddle-gpu version, which
|
||||
# may copy CPU tensor to GPU even if users want to use
|
||||
# CPU tensor operation, so we add CPUPlace guard here
|
||||
# to make sure tensor will be operated only on CPU
|
||||
with paddle.base.dygraph.guard(place=paddle.CPUPlace()):
|
||||
batch = fetcher.fetch(indices)
|
||||
except Exception as e:
|
||||
if (
|
||||
isinstance(e, StopIteration)
|
||||
and dataset_kind == _DatasetKind.ITER
|
||||
):
|
||||
out_queue.put(_IterableDatasetStopIteration(worker_id))
|
||||
iterator_drained = True
|
||||
else:
|
||||
out_queue.put((idx, _WorkerException(worker_id), None))
|
||||
else:
|
||||
if isinstance(batch, _WorkerException):
|
||||
out_queue.put((idx, batch, None))
|
||||
batch, structure = _flatten_batch(batch)
|
||||
if use_shared_memory:
|
||||
|
||||
def numpy2lodtensor(arr):
|
||||
lodtensor = core.DenseTensor()
|
||||
lodtensor.set(arr, core.CPUPlace())
|
||||
return lodtensor
|
||||
|
||||
tensor_list = [
|
||||
(
|
||||
numpy2lodtensor(b)
|
||||
if isinstance(b, np.ndarray)
|
||||
else b.get_tensor()
|
||||
)
|
||||
for b in batch
|
||||
]
|
||||
out_queue.put((idx, tensor_list, structure))
|
||||
else:
|
||||
out_queue.put((idx, batch, structure))
|
||||
except KeyboardInterrupt:
|
||||
# NOTE: Main process will raise KeyboardInterrupt anyways, ignore it in child process
|
||||
pass
|
||||
except:
|
||||
raise
|
||||
finally:
|
||||
if use_shared_memory:
|
||||
_cleanup_mmap()
|
||||
if done_event.is_set():
|
||||
out_queue.cancel_join_thread()
|
||||
out_queue.close()
|
||||
Reference in New Issue
Block a user