Files
2026-07-13 12:40:42 +08:00

766 lines
22 KiB
Python

# Copyright (c) 2016 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 itertools
import logging
import multiprocessing
import random
import sys
import warnings
from collections.abc import Callable, Generator
from itertools import zip_longest
from queue import Queue
from threading import Thread
from typing import (
TYPE_CHECKING,
Any,
TypeAlias,
TypedDict,
TypeVar,
overload,
)
from typing_extensions import NotRequired, Unpack
from paddle.base.reader import QUEUE_GET_TIMEOUT
if TYPE_CHECKING:
from collections.abc import Sequence
class _ComposeOptions(TypedDict):
check_alignment: NotRequired[bool]
__all__ = []
# On macOS, Python uses the 'spawn' start method by default in multiprocessing.
# Paddle is currently unable to solve this, so forces the process to start using
# the 'fork' start method.
#
# TODO: This solution is not good, because the fork start method could lead to
# crashes of the subprocess. Figure out how to make 'spawn' work.
#
# For more details, please refer to
# https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
# https://bugs.python.org/issue33725
if sys.platform == 'darwin':
fork_context = multiprocessing.get_context('fork')
else:
fork_context = multiprocessing
_T = TypeVar('_T')
_T1 = TypeVar('_T1')
_T2 = TypeVar('_T2')
_T3 = TypeVar('_T3')
_T4 = TypeVar('_T4')
_U = TypeVar('_U')
_Reader: TypeAlias = Callable[[], Generator[_T, None, None]]
def cache(reader: _Reader[_T]) -> _Reader[_T]:
"""
Cache the reader data into memory.
Be careful that this method may take long time to process,
and consume lots of memory. :code:`reader()` would only
call once.
Args:
reader (generator): a reader object which yields
data each time.
Returns:
generator: a decorated reader object which yields data from cached memory.
Examples:
.. code-block:: pycon
>>> import paddle
>>> def reader():
... for i in range(3):
... yield i
>>> # All data is cached into memory
>>> cached_reader = paddle.base.io.cache(reader)
>>> for i in cached_reader():
... print(i)
0
1
2
"""
all_data = tuple(reader())
def __impl__() -> Generator[_T, None, None]:
yield from all_data
return __impl__
# A temporary solution like builtin map function.
# `Map` maybe the final solution in the future.
# See https://github.com/python/typing/issues/1383
@overload
def map_readers(
func: Callable[[_T1], _U], reader1: _Reader[_T1], /
) -> _Reader[_U]: ...
@overload
def map_readers(
func: Callable[[_T1, _T2], _U],
reader1: _Reader[_T1],
reader2: _Reader[_T2],
/,
) -> _Reader[_U]: ...
@overload
def map_readers(
func: Callable[[_T1, _T2, _T3], _U],
reader1: _Reader[_T1],
reader2: _Reader[_T2],
reader3: _Reader[_T3],
/,
) -> _Reader[_U]: ...
@overload
def map_readers(
func: Callable[[_T1, _T2, _T3, _T4], _U],
reader1: _Reader[_T1],
reader2: _Reader[_T2],
reader3: _Reader[_T3],
reader4: _Reader[_T4],
/,
) -> _Reader[_U]: ...
@overload
def map_readers(
func: Callable[..., _U], *readers: _Reader[Any]
) -> _Reader[_U]: ...
def map_readers(func, *readers):
"""
Creates a data reader that outputs return value of function using
output of each data reader as arguments.
If input readers output the following data entries: 2 3,
and the input func is mul(x, y),
the output of the resulted reader will be 6.
Args:
func: a function to read data and compute result, the output of this function
will be set as the output of the resulted data reader.
readers (Reader|list of Reader): list of readers whose outputs will be used as arguments of func.
Returns:
the resulted data reader (Reader)
Examples:
.. code-block:: pycon
>>> import paddle.reader
>>> d = {"h": 0, "i": 1}
>>> def func(x):
... return d[x]
>>> def reader():
... yield "h"
... yield "i"
>>> map_reader_result = paddle.reader.map_readers(func, reader)
"""
def reader():
rs = []
for r in readers:
rs.append(r())
yield from map(func, *rs)
return reader
def shuffle(reader: _Reader[_T], buf_size: int) -> _Reader[_T]:
"""
This API creates a decorated reader that outputs the shuffled data.
The output data from the origin reader will be saved into a buffer,
and then shuffle the data. The size of buffer is determined by argument buf_size.
Args:
reader(callable): the original reader whose data will be shuffled.
buf_size(int): the size of shuffled buffer.
Returns:
callable: a decorated reader.
Examples:
.. code-block:: pycon
>>> # doctest: +SKIP('outputs are 0~4 unordered arrangement')
>>> def reader():
... for i in range(5):
... yield i
>>> shuffled_reader = paddle.reader.decorator.shuffle(reader, 3)
>>> for e in shuffled_reader():
... print(e)
>>> # outputs are 0~4 unordered arrangement
"""
def data_reader() -> Generator[_T, None, None]:
buf = []
for e in reader():
buf.append(e)
if len(buf) >= buf_size:
random.shuffle(buf)
for b in buf:
yield b
buf = []
if len(buf) > 0:
random.shuffle(buf)
for b in buf:
yield b
return data_reader
def chain(*readers: _Reader[_T]) -> _Reader[_T]:
"""
Use the input data readers to create a chained data reader. The new created reader
chains the outputs of input readers together as its output, and it do not change
the format of the outputs.
**Note**:
``paddle.reader.chain`` is the alias of ``paddle.base.io.chain``, and
``paddle.base.io.chain`` is recommended to use.
For example, if three input readers' outputs are as follows:
[0, 0, 0],
[10, 10, 10],
[20, 20, 20].
The chained reader will output:
[0, 0, 0], [10, 10, 10], [20, 20, 20].
Args:
readers(list): input data readers.
Returns:
callable: the new chained data reader.
Examples:
.. code-block:: pycon
>>> import paddle
>>> def reader_creator_3(start):
... def reader():
... for i in range(start, start + 3):
... yield [i, i, i]
...
... return reader
>>> c = paddle.reader.chain(reader_creator_3(0), reader_creator_3(10), reader_creator_3(20))
>>> for e in c():
... print(e)
[0, 0, 0]
[1, 1, 1]
[2, 2, 2]
[10, 10, 10]
[11, 11, 11]
[12, 12, 12]
[20, 20, 20]
[21, 21, 21]
[22, 22, 22]
"""
def reader() -> Generator[_T, None, None]:
rs: list[Generator[_T, None, None]] = []
for r in readers:
rs.append(r())
yield from itertools.chain(*rs)
return reader
class ComposeNotAligned(ValueError):
pass
def compose(
*readers: _Reader[Any], **kwargs: Unpack[_ComposeOptions]
) -> _Reader[Any]:
"""
Creates a data reader whose output is the combination of input readers.
If input readers output following data entries:
(1, 2) 3 (4, 5)
The composed reader will output:
(1, 2, 3, 4, 5)
Args:
readers (Reader|list of Reader): readers that will be composed together.
check_alignment(bool, optional): Indicates whether the input readers are checked for
alignment. If True, whether input readers are aligned
correctly will be checked, else alignment will not be checkout and trailing outputs
will be discarded. Defaults to True.
Returns:
the new data reader (Reader).
Examples:
.. code-block:: pycon
>>> def reader_creator_10(dur):
... def reader():
... for i in range(10):
... yield i
...
... return reader
>>> reader = paddle.reader.decorator.compose(reader_creator_10(0), reader_creator_10(0))
"""
check_alignment = kwargs.pop('check_alignment', True)
def make_tuple(x):
if isinstance(x, tuple):
return x
else:
return (x,)
def reader():
rs = []
for r in readers:
rs.append(r())
if not check_alignment:
for outputs in zip(*rs):
yield sum(list(map(make_tuple, outputs)), ())
else:
for outputs in zip_longest(*rs):
for o in outputs:
if o is None:
# None will be not be present if compose is aligned
raise ComposeNotAligned(
"outputs of readers are not aligned."
)
yield sum(list(map(make_tuple, outputs)), ())
return reader
def buffered(reader: _Reader[_T], size: int) -> _Reader[_T]:
"""
Creates a buffered data reader.
The buffered data reader will read and save data entries into a
buffer. Reading from the buffered data reader will proceed as long
as the buffer is not empty.
Args:
reader(generator): the data reader to read from.
size(int): max buffer size.
Returns:
generator: the buffered data reader.
Examples:
.. code-block:: pycon
>>> import paddle
>>> def reader():
... for i in range(3):
... yield i
>>> # Create a buffered reader, and the buffer size is 2.
>>> buffered_reader = paddle.reader.decorator.buffered(reader, 2)
>>> # Output: 0 1 2
>>> for i in buffered_reader():
... print(i)
0
1
2
"""
class EndSignal:
pass
end = EndSignal()
def read_worker(r, q):
for d in r:
q.put(d)
q.put(end)
def data_reader():
r = reader()
q = Queue(maxsize=size)
t = Thread(
target=read_worker,
args=(r, q),
)
t.daemon = True
t.start()
e = q.get()
while e != end:
yield e
e = q.get()
return data_reader
def firstn(reader: _Reader[_T], n: int) -> _Reader[_T]:
"""
This API creates a decorated reader, and limits the max number of
samples that reader could return.
Args:
reader(callable): the input reader.
n(int): the max number of samples in the reader.
Returns:
callable: the decorated reader.
Examples:
.. code-block:: pycon
>>> def reader():
... for i in range(100):
... yield i
>>> firstn_reader = paddle.reader.decorator.firstn(reader, 5)
>>> for e in firstn_reader():
... print(e)
0
1
2
3
4
"""
# TODO(yuyang18): Check if just drop the reader, could clean the opened
# resource or not?
def firstn_reader():
for i, item in enumerate(reader()):
if i == n:
break
yield item
return firstn_reader
class XmapEndSignal:
pass
def xmap_readers(
mapper: Callable[[_T], _U],
reader: _Reader[_T],
process_num: int,
buffer_size: int,
order: bool = False,
) -> _Reader[_U]:
"""
Use multi-threads to map samples from reader by a mapper defined by user.
Args:
mapper (callable): a function to map the data from reader.
reader (callable): a data reader which yields the data.
process_num (int): thread number to handle original sample.
buffer_size (int): size of the queue to read data in.
order (bool): whether to keep the data order from original reader.
Default False.
Returns:
callable: a decorated reader with data mapping.
"""
end = XmapEndSignal()
# define a worker to read samples from reader to in_queue
def read_worker(reader, in_queue):
for i in reader():
in_queue.put(i)
in_queue.put(end)
# define a worker to read samples from reader to in_queue with order flag
def order_read_worker(reader, in_queue):
in_order = 0
for i in reader():
in_queue.put((in_order, i))
in_order += 1
in_queue.put(end)
# define a worker to handle samples from in_queue by mapper
# and put mapped samples into out_queue
def handle_worker(in_queue, out_queue, mapper):
sample = in_queue.get()
while not isinstance(sample, XmapEndSignal):
r = mapper(sample)
out_queue.put(r)
sample = in_queue.get()
in_queue.put(end)
out_queue.put(end)
# define a worker to handle samples from in_queue by mapper
# and put mapped samples into out_queue by order
def order_handle_worker(in_queue, out_queue, mapper, out_order):
ins = in_queue.get()
while not isinstance(ins, XmapEndSignal):
order, sample = ins
r = mapper(sample)
while order != out_order[0]:
pass
out_queue.put(r)
out_order[0] += 1
ins = in_queue.get()
in_queue.put(end)
out_queue.put(end)
def xreader():
in_queue = Queue(buffer_size)
out_queue = Queue(buffer_size)
out_order = [0]
# start a read worker in a thread
target = order_read_worker if order else read_worker
t = Thread(target=target, args=(reader, in_queue))
t.daemon = True
t.start()
# start several handle_workers
target = order_handle_worker if order else handle_worker
args = (
(in_queue, out_queue, mapper, out_order)
if order
else (in_queue, out_queue, mapper)
)
workers = []
for i in range(process_num):
worker = Thread(target=target, args=args)
worker.daemon = True
workers.append(worker)
for w in workers:
w.start()
sample = out_queue.get()
while not isinstance(sample, XmapEndSignal):
yield sample
sample = out_queue.get()
finish = 1
while finish < process_num:
sample = out_queue.get()
if isinstance(sample, XmapEndSignal):
finish += 1
else:
yield sample
return xreader
def multiprocess_reader(
readers: Sequence[_Reader[_T]],
use_pipe: bool = True,
queue_size: int = 1000,
) -> _Reader[list[_T]]:
"""
This API use python ``multiprocessing`` to read data from ``readers`` parallelly,
and then ``multiprocess.Queue`` or ``multiprocess.Pipe`` is used to merge
these data. A separate process will be created for each reader in the
``readers`` list, please guarantee every reader can work independently
to avoid conflicts in parallel environment.
``Multiprocess.Queue`` require the rw access right to /dev/shm, and it's not supported
in some platforms.
Parameters:
readers (list( ``generator`` ) | tuple( ``generator`` )): a python ``generator`` list
used to read input data
use_pipe (bool, optional): control the inner API used to implement the multi-processing,
default True - use ``multiprocess.Pipe`` which is recommended
queue_size (int, optional): only useful when ``use_pipe`` is False - ``multiprocess.Queue``
is used, default 1000. Increase this value can speed up the data reading, and more memory
will be consumed.
Returns:
``generator``: a new reader which can be run parallelly
Example:
.. code-block:: pycon
>>> import paddle
>>> import numpy as np
>>> sample_files = ['sample_file_1', 'sample_file_2']
>>> def fake_input_files():
... with open(sample_files[0], 'wb') as f:
... np.savez(f, a=np.array([1, 2]), b=np.array([3, 4]), c=np.array([5, 6]), d=np.array([7, 8]))
... with open(sample_files[1], 'wb') as f:
... np.savez(f, a=np.array([9, 10]), b=np.array([11, 12]), c=np.array([13, 14]))
>>> def generate_reader(file_name):
... # load data file
... def _impl():
... data = np.load(file_name)
... for item in sorted(data.files):
... yield (data[item],)
...
... return _impl
>>> if __name__ == '__main__':
... # generate sample input files
... fake_input_files()
...
... with base.program_guard(base.Program(), base.Program()):
... place = base.CPUPlace()
... # the 1st 2 is batch size
...
... image = paddle.static.data(name='image', dtype='int64', shape=[2, 1, 2])
... paddle.static.Print(image)
... # print detailed tensor info of image variable
...
... reader = base.io.PyReader(feed_list=[image], capacity=2)
...
... decorated_reader = paddle.reader.multiprocess_reader(
... [generate_reader(sample_files[0]), generate_reader(sample_files[1])], False
... )
...
... reader.decorate_sample_generator(decorated_reader, batch_size=2, places=[place])
...
... exe = base.Executor(place)
... exe.run(base.default_startup_program())
...
... for data in reader():
... res = exe.run(feed=data, fetch_list=[image])
... print(res[0])
[[[1 2]], [[3 4]]]
[[[5 6]], [[7 8]]]
[[[9 10]], [[11 12]]]
"""
if sys.platform == 'win32':
raise NotImplementedError(
"The multiprocess_reader method is not supported on windows."
)
# ujson is ultra fast json encoder and decoder written in pure C with bindings for Python 3.6+.
try:
import ujson as json
except Exception as e:
warnings.warn(
"The `ujson` module is not found, use the `json` module, `ujson` encodes and decodes faster, "
"you can install `ujson` through `pip install ujson`."
)
import json
assert isinstance(readers, (list, tuple)) and len(readers) > 0, (
"`readers` must be list or tuple."
)
def _read_into_queue(reader, queue):
try:
for sample in reader():
if sample is None:
raise ValueError("sample has None")
queue.put(sample)
queue.put(None)
except Exception as e:
queue.put("")
raise e
def queue_reader():
queue = fork_context.Queue(queue_size)
for reader in readers:
p = fork_context.Process(
target=_read_into_queue, args=(reader, queue)
)
p.start()
reader_num = len(readers)
finish_num = 0
while finish_num < reader_num:
try:
sample = queue.get(timeout=QUEUE_GET_TIMEOUT)
except Exception as e:
logging.error(
"multiprocess_reader failed to get data from the multiprocessing.Queue."
)
raise e
if sample is None:
finish_num += 1
elif sample == "":
raise ValueError(
"multiprocess_reader failed to put data into the multiprocessing.Queue."
)
else:
yield sample
def _read_into_pipe(reader, conn):
try:
for sample in reader():
if sample is None:
raise ValueError("sample has None!")
conn.send(json.dumps(sample))
conn.send(json.dumps(None))
conn.close()
except Exception as e:
conn.send(json.dumps(""))
conn.close()
raise e
def pipe_reader():
conns = []
for reader in readers:
parent_conn, child_conn = fork_context.Pipe()
conns.append(parent_conn)
p = fork_context.Process(
target=_read_into_pipe, args=(reader, child_conn)
)
p.start()
reader_num = len(readers)
finish_num = 0
conn_to_remove = []
while finish_num < reader_num:
for conn in conn_to_remove:
conns.remove(conn)
conn_to_remove = []
for conn in conns:
sample = json.loads(conn.recv())
if sample is None:
finish_num += 1
conn.close()
conn_to_remove.append(conn)
elif sample == "":
conn.close()
conn_to_remove.append(conn)
raise ValueError(
"multiprocess_reader failed to send data into the multiprocessing.Pipe."
)
else:
yield sample
if use_pipe:
return pipe_reader
else:
return queue_reader