base class for defining a parallel dataset using python code every pydataset must implement the getitem and the len methods if you want to modify your dataset between epochs you may additionally implement onepochend the getitem method should return a complete batch not a single sample and the len method should return the number of batches in the dataset rather than the number of samples args workers number of workers to use in multithreading or multiprocessing usemultiprocessing whether to use python multiprocessing for parallelism setting this to true means that your dataset will be replicated in multiple forked processes this is necessary to gain computelevel rather than io level benefits from parallelism however it can only be set to true if your dataset can be safely pickled maxqueuesize maximum number of batches to keep in the queue when iterating over the dataset in a multithreaded or multipricessed setting reduce this value to reduce the cpu memory consumption of your dataset defaults to 10 notes pydataset is a safer way to do multiprocessing this structure guarantees that the model will only train once on each sample per epoch which is not the case with python generators the arguments workers usemultiprocessing and maxqueuesize exist to configure how fit uses parallelism to iterate over the dataset they are not being used by the pydataset class directly when you are manually iterating over a pydataset no parallelism is applied example python from skimage io import imread from skimage transform import resize import numpy as np import math here xset is list of path to the images and yset are the associated classes class cifar10pydatasetkeras utils pydataset def initself xset yset batchsize kwargs super initkwargs self x self y xset yset self batchsize batchsize def lenself return number of batches return math ceillenself x self batchsize def getitemself idx return x y for batch idx low idx self batchsize cap upper bound at array length the last batch may be smaller if the total number of items is not a multiple of batch size high minlow self batchsize lenself x batchx self xlow high batchy self ylow high return np array resizeimreadfilename 200 200 for filename in batchx np arraybatchy gets batch at position index args index position of the batch in the pydataset returns a batch number of batch in the pydataset returns the number of batches in the pydataset method called at the end of every epoch pass def iterself adapter for keras utils pydataset instances def init self x classweightnone shufflefalse self pydataset x self classweight classweight self enqueuer none self shuffle shuffle self outputsignature none def settfoutputsignatureself from keras utils moduleutils import tensorflow as tf def gettensorspecx shape x shape if lenshape 1 raise valueerror the arrays returned by pydataset getitem must be at least rank 1 received fx of rank lenx shape shape listshape shape0 none the batch size is not guaranteed to be static return tf tensorspecshapeshape dtypex dtype name grab the first example batch self pydataset0 run checks on it and format it batch self standardizebatchbatch self outputsignature tree mapstructuregettensorspec batch def standardizebatchself batch if isinstancebatch dict return batch if isinstancebatch np ndarray batch batch if isinstancebatch list batch tuplebatch if not isinstancebatch tuple or lenbatch not in 1 2 3 raise valueerror pydataset getitem must return a tuple or a dict if a tuple it must be ordered either input or inputs targets or inputs targets sampleweights freceived strbatch 100 of type typebatch if self classweight is not none if lenbatch 3 raise valueerror you cannot specify classweight and sampleweight at the same time if lenbatch 2 sw dataadapterutils classweighttosampleweights batch1 self classweight batch batch sw return batch def makemultiprocessedgeneratorfnself workers self pydataset workers usemultiprocessing self pydataset usemultiprocessing if workers 1 or workers 0 and usemultiprocessing def generatorfn self enqueuer orderedenqueuer self pydataset usemultiprocessingusemultiprocessing shuffleself shuffle self enqueuer start workersworkers maxqueuesizeself pydataset maxqueuesize return self enqueuer get else def generatorfn order rangelenself pydataset if self shuffle match the shuffle convention in orderedenqueuer order listorder random shuffleorder for i in order yield self pydataseti return generatorfn def getnumpyiteratorself genfn self makemultiprocessedgeneratorfn for i batch in enumerategenfn batch self standardizebatchbatch yield batch if i lenself pydataset 1 and self enqueuer self enqueuer stop def gettfdatasetself from keras utils moduleutils import tensorflow as tf if self outputsignature is none self settfoutputsignature ds tf data dataset fromgenerator self getnumpyiterator outputsignatureself outputsignature if self shuffle ds ds shuffle8 ds ds prefetchtf data autotune return ds def onepochendself if self enqueuer self enqueuer stop self pydataset onepochend property def numbatchesself return lenself pydataset property def batchsizeself return none global variables to be shared across processes sharedsequences we use a value to provide unique id to different processes sequencecounter none because multiprocessing pools are inherently unsafe starting from a clean state can be essential to avoiding deadlocks in order to accomplish this we need to be able to check on the status of pools that we create datapools weakref weakset workeridqueue none only created if needed forcethreadpool false def getpoolclassusemultiprocessing global forcethreadpool if not usemultiprocessing or forcethreadpool return multiprocessing dummy pool threadpool return multiprocessing pool def getworkeridqueue get the value from the pydataset uid at index i to allow multiple pydatasets to be used at the same time we use uid to get a specific one a single pydataset would cause the validation to overwrite the training pydataset args uid int pydataset identifier i index returns the value at index i base class to enqueue inputs the task of an enqueuer is to use parallelism to speed up preprocessing this is done with processes or threads example python enqueuer pydatasetenqueuer enqueuer start datas enqueuer get for data in datas use the inputs training evaluating predicting stop sometime enqueuer stop the enqueuer get should be an infinite stream of data in this case the os does not allow us to use multiprocessing we resort to an int for enqueuer indexing doing multiprocessing value x is not processsafe starts the handler s workers args workers number of workers maxqueuesize queue size when full workers could block on put we do not need the init since it s threads sends current iterable to all workers for new processes that may spawn sharedsequencesself uid self pydataset def stopself timeoutnone if not self isrunning return self stopsignal set with self queue mutex self queue queue clear self queue unfinishedtasks 0 self queue notfull notify self runthread jointimeout sharedsequencesself uid none def delself if self isrunning self stop def runself gets the pool initializer for multiprocessing args workers number of workers returns function a function to initialize the pool creates a generator to extract data from the queue skip the data if it is none returns generator yielding tuples inputs targets or inputs targets sampleweights builds a enqueuer from a pydataset args pydataset a keras utils pydataset object usemultiprocessing use multiprocessing if true otherwise threading shuffle whether to shuffle the data at the beginning of each epoch gets the pool initializer for multiprocessing args workers number of workers returns function a function to initialize the pool wait for the queue to be empty while true time sleep0 1 if self queue unfinishedtasks 0 or self stopsignal isset return def runself done with the current epoch waiting for the final batches we re done call the internal on epoch end creates a generator to extract data from the queue skip the data if it is none yields the next element in the queue i e a tuple inputs targets or inputs targets sampleweights initializer function for pool workers args gens state which should be made available to worker processes randomseed an optional value with which to seed child processes idqueue a multiprocessing queue of worker ids this is used to indicate that a worker process was keras name isn t used for anything but setting a more descriptive name is helpful when diagnosing orphaned processes if a worker dies during init the pool will just create a replacement base class for defining a parallel dataset using python code every pydataset must implement the __getitem__ and the __len__ methods if you want to modify your dataset between epochs you may additionally implement on_epoch_end the __getitem__ method should return a complete batch not a single sample and the __len__ method should return the number of batches in the dataset rather than the number of samples args workers number of workers to use in multithreading or multiprocessing use_multiprocessing whether to use python multiprocessing for parallelism setting this to true means that your dataset will be replicated in multiple forked processes this is necessary to gain compute level rather than i o level benefits from parallelism however it can only be set to true if your dataset can be safely pickled max_queue_size maximum number of batches to keep in the queue when iterating over the dataset in a multithreaded or multipricessed setting reduce this value to reduce the cpu memory consumption of your dataset defaults to 10 notes pydataset is a safer way to do multiprocessing this structure guarantees that the model will only train once on each sample per epoch which is not the case with python generators the arguments workers use_multiprocessing and max_queue_size exist to configure how fit uses parallelism to iterate over the dataset they are not being used by the pydataset class directly when you are manually iterating over a pydataset no parallelism is applied example python from skimage io import imread from skimage transform import resize import numpy as np import math here x_set is list of path to the images and y_set are the associated classes class cifar10pydataset keras utils pydataset def __init__ self x_set y_set batch_size kwargs super __init__ kwargs self x self y x_set y_set self batch_size batch_size def __len__ self return number of batches return math ceil len self x self batch_size def __getitem__ self idx return x y for batch idx low idx self batch_size cap upper bound at array length the last batch may be smaller if the total number of items is not a multiple of batch size high min low self batch_size len self x batch_x self x low high batch_y self y low high return np array resize imread file_name 200 200 for file_name in batch_x np array batch_y gets batch at position index args index position of the batch in the pydataset returns a batch number of batch in the pydataset returns the number of batches in the pydataset method called at the end of every epoch create a generator that iterate over the pydataset adapter for keras utils pydataset instances the batch size is not guaranteed to be static grab the first example run checks on it and format it match the shuffle convention in orderedenqueuer global variables to be shared across processes we use a value to provide unique id to different processes because multiprocessing pools are inherently unsafe starting from a clean state can be essential to avoiding deadlocks in order to accomplish this we need to be able to check on the status of pools that we create only created if needed threadpool lazily create the queue to track worker ids get the value from the pydataset uid at index i to allow multiple pydatasets to be used at the same time we use uid to get a specific one a single pydataset would cause the validation to overwrite the training pydataset args uid int pydataset identifier i index returns the value at index i base class to enqueue inputs the task of an enqueuer is to use parallelism to speed up preprocessing this is done with processes or threads example python enqueuer pydatasetenqueuer enqueuer start datas enqueuer get for data in datas use the inputs training evaluating predicting stop sometime enqueuer stop the enqueuer get should be an infinite stream of data in this case the os does not allow us to use multiprocessing we resort to an int for enqueuer indexing doing multiprocessing value x is not process safe starts the handler s workers args workers number of workers max_queue_size queue size when full workers could block on put we do not need the init since it s threads sends current iterable to all workers for new processes that may spawn stops running threads and wait for them to exit if necessary should be called by the same thread which called start args timeout maximum time to wait on thread join submits request to the executor and queue the future objects gets the pool initializer for multiprocessing args workers number of workers returns function a function to initialize the pool creates a generator to extract data from the queue skip the data if it is none returns generator yielding tuples inputs targets or inputs targets sample_weights builds a enqueuer from a pydataset args py_dataset a keras utils pydataset object use_multiprocessing use multiprocessing if true otherwise threading shuffle whether to shuffle the data at the beginning of each epoch gets the pool initializer for multiprocessing args workers number of workers returns function a function to initialize the pool wait for the queue to be empty submits request to the executor and queue the future objects share the initial py_dataset done with the current epoch waiting for the final batches we re done call the internal on epoch end update the pool creates a generator to extract data from the queue skip the data if it is none yields the next element in the queue i e a tuple inputs targets or inputs targets sample_weights initializer function for pool workers args gens state which should be made available to worker processes random_seed an optional value with which to seed child processes id_queue a multiprocessing queue of worker ids this is used to indicate that a worker process was keras name isn t used for anything but setting a more descriptive name is helpful when diagnosing orphaned processes if a worker dies during init the pool will just create a replacement
import multiprocessing
import queue
import random
import threading
import time
import warnings
import weakref
from contextlib import closing

import numpy as np
import tree

from keras.api_export import keras_export
from keras.trainers.data_adapters import data_adapter_utils
from keras.trainers.data_adapters.data_adapter import DataAdapter


@keras_export(["keras.utils.PyDataset", "keras.utils.Sequence"])
class PyDataset:
    

    def __init__(self, workers=1, use_multiprocessing=False, max_queue_size=10):
        self._workers = workers
        self._use_multiprocessing = use_multiprocessing
        self._max_queue_size = max_queue_size

    def _warn_if_super_not_called(self):
        warn = False
        if not hasattr(self, "_workers"):
            self._workers = 1
            warn = True
        if not hasattr(self, "_use_multiprocessing"):
            self._use_multiprocessing = False
            warn = True
        if not hasattr(self, "_max_queue_size"):
            self._max_queue_size = 10
            warn = True
        if warn:
            warnings.warn(
                "Your `PyDataset` class should call "
                "`super().__init__(**kwargs)` in its constructor. "
                "`**kwargs` can include `workers`, "
                "`use_multiprocessing`, `max_queue_size`. Do not pass "
                "these arguments to `fit()`, as they will be ignored.",
                stacklevel=2,
            )

    @property
    def workers(self):
        self._warn_if_super_not_called()
        return self._workers

    @workers.setter
    def workers(self, value):
        self._workers = value

    @property
    def use_multiprocessing(self):
        self._warn_if_super_not_called()
        return self._use_multiprocessing

    @use_multiprocessing.setter
    def use_multiprocessing(self, value):
        self._use_multiprocessing = value

    @property
    def max_queue_size(self):
        self._warn_if_super_not_called()
        return self._max_queue_size

    @max_queue_size.setter
    def max_queue_size(self, value):
        self._max_queue_size = value

    def __getitem__(self, index):
        
        raise NotImplementedError

    def __len__(self):
        
        raise NotImplementedError

    def on_epoch_end(self):
        
        pass

    def __iter__(self):
        
        for i in range(len(self)):
            yield self[i]


class PyDatasetAdapter(DataAdapter):
    

    def __init__(
        self,
        x,
        class_weight=None,
        shuffle=False,
    ):
        self.py_dataset = x
        self.class_weight = class_weight
        self.enqueuer = None
        self.shuffle = shuffle
        self._output_signature = None

    def _set_tf_output_signature(self):
        from keras.utils.module_utils import tensorflow as tf

        def get_tensor_spec(x):
            shape = x.shape
            if len(shape) < 1:
                raise ValueError(
                    "The arrays returned by PyDataset.__getitem__() "
                    "must be at least rank 1. Received: "
                    f"{x} of rank {len(x.shape)}"
                )
            shape = list(shape)
            shape[0] = None  
            return tf.TensorSpec(shape=shape, dtype=x.dtype.name)

        
        batch = self.py_dataset[0]
        
        batch = self._standardize_batch(batch)
        self._output_signature = tree.map_structure(get_tensor_spec, batch)

    def _standardize_batch(self, batch):
        if isinstance(batch, dict):
            return batch
        if isinstance(batch, np.ndarray):
            batch = (batch,)
        if isinstance(batch, list):
            batch = tuple(batch)
        if not isinstance(batch, tuple) or len(batch) not in {1, 2, 3}:
            raise ValueError(
                "PyDataset.__getitem__() must return a tuple or a dict. "
                "If a tuple, it must be ordered either "
                "(input,) or (inputs, targets) or "
                "(inputs, targets, sample_weights). "
                f"Received: {str(batch)[:100]}... of type {type(batch)}"
            )
        if self.class_weight is not None:
            if len(batch) == 3:
                raise ValueError(
                    "You cannot specify `class_weight` "
                    "and `sample_weight` at the same time."
                )
            if len(batch) == 2:
                sw = data_adapter_utils.class_weight_to_sample_weights(
                    batch[1], self.class_weight
                )
                batch = batch + (sw,)
        return batch

    def _make_multiprocessed_generator_fn(self):
        workers = self.py_dataset.workers
        use_multiprocessing = self.py_dataset.use_multiprocessing
        if workers > 1 or (workers > 0 and use_multiprocessing):

            def generator_fn():
                self.enqueuer = OrderedEnqueuer(
                    self.py_dataset,
                    use_multiprocessing=use_multiprocessing,
                    shuffle=self.shuffle,
                )
                self.enqueuer.start(
                    workers=workers,
                    max_queue_size=self.py_dataset.max_queue_size,
                )
                return self.enqueuer.get()

        else:

            def generator_fn():
                order = range(len(self.py_dataset))
                if self.shuffle:
                    
                    order = list(order)
                    random.shuffle(order)

                for i in order:
                    yield self.py_dataset[i]

        return generator_fn

    def get_numpy_iterator(self):
        gen_fn = self._make_multiprocessed_generator_fn()
        for i, batch in enumerate(gen_fn()):
            batch = self._standardize_batch(batch)
            yield batch
            if i >= len(self.py_dataset) - 1 and self.enqueuer:
                self.enqueuer.stop()

    def get_tf_dataset(self):
        from keras.utils.module_utils import tensorflow as tf

        if self._output_signature is None:
            self._set_tf_output_signature()

        ds = tf.data.Dataset.from_generator(
            self.get_numpy_iterator,
            output_signature=self._output_signature,
        )
        if self.shuffle:
            ds = ds.shuffle(8)
        ds = ds.prefetch(tf.data.AUTOTUNE)
        return ds

    def on_epoch_end(self):
        if self.enqueuer:
            self.enqueuer.stop()
        self.py_dataset.on_epoch_end()

    @property
    def num_batches(self):
        return len(self.py_dataset)

    @property
    def batch_size(self):
        return None



_SHARED_SEQUENCES = {}

_SEQUENCE_COUNTER = None





_DATA_POOLS = weakref.WeakSet()
_WORKER_ID_QUEUE = None  
_FORCE_THREADPOOL = False


def get_pool_class(use_multiprocessing):
    global _FORCE_THREADPOOL
    if not use_multiprocessing or _FORCE_THREADPOOL:
        return multiprocessing.dummy.Pool  
    return multiprocessing.Pool


def get_worker_id_queue():
    
    global _WORKER_ID_QUEUE
    if _WORKER_ID_QUEUE is None:
        _WORKER_ID_QUEUE = multiprocessing.Queue()
    return _WORKER_ID_QUEUE


def init_pool(seqs):
    global _SHARED_SEQUENCES
    _SHARED_SEQUENCES = seqs


def get_index(uid, i):
    
    return _SHARED_SEQUENCES[uid][i]


class PyDatasetEnqueuer:
    

    def __init__(self, py_dataset, use_multiprocessing=False):
        self.py_dataset = py_dataset
        self.use_multiprocessing = use_multiprocessing

        global _SEQUENCE_COUNTER
        if _SEQUENCE_COUNTER is None:
            try:
                _SEQUENCE_COUNTER = multiprocessing.Value("i", 0)
            except OSError:
                
                
                
                _SEQUENCE_COUNTER = 0

        if isinstance(_SEQUENCE_COUNTER, int):
            self.uid = _SEQUENCE_COUNTER
            _SEQUENCE_COUNTER += 1
        else:
            
            with _SEQUENCE_COUNTER.get_lock():
                self.uid = _SEQUENCE_COUNTER.value
                _SEQUENCE_COUNTER.value += 1

        self.workers = 0
        self.executor_fn = None
        self.queue = None
        self.run_thread = None
        self.stop_signal = None

    def is_running(self):
        return self.stop_signal is not None and not self.stop_signal.is_set()

    def start(self, workers=1, max_queue_size=10):
        
        if self.use_multiprocessing:
            self.executor_fn = self._get_executor_init(workers)
        else:
            
            self.executor_fn = lambda _: get_pool_class(False)(workers)
        self.workers = workers
        self.queue = queue.Queue(max_queue_size)
        self.stop_signal = threading.Event()
        self.run_thread = threading.Thread(target=self._run)
        self.run_thread.daemon = True
        self.run_thread.start()

    def _send_py_dataset(self):
        
        
        _SHARED_SEQUENCES[self.uid] = self.py_dataset

    def stop(self, timeout=None):
        
        if not self.is_running():
            return
        self.stop_signal.set()
        with self.queue.mutex:
            self.queue.queue.clear()
            self.queue.unfinished_tasks = 0
            self.queue.not_full.notify()
        self.run_thread.join(timeout)
        _SHARED_SEQUENCES[self.uid] = None

    def __del__(self):
        if self.is_running():
            self.stop()

    def _run(self):
        
        raise NotImplementedError

    def _get_executor_init(self, workers):
        
        raise NotImplementedError

    def get(self):
        
        raise NotImplementedError


class OrderedEnqueuer(PyDatasetEnqueuer):
    

    def __init__(self, py_dataset, use_multiprocessing=False, shuffle=False):
        super().__init__(py_dataset, use_multiprocessing)
        self.shuffle = shuffle

    def _get_executor_init(self, workers):
        

        def pool_fn(seqs):
            pool = get_pool_class(True)(
                workers,
                initializer=init_pool_generator,
                initargs=(seqs, None, get_worker_id_queue()),
            )
            _DATA_POOLS.add(pool)
            return pool

        return pool_fn

    def _wait_queue(self):
        
        while True:
            time.sleep(0.1)
            if self.queue.unfinished_tasks == 0 or self.stop_signal.is_set():
                return

    def _run(self):
        
        indices = list(range(len(self.py_dataset)))
        if self.shuffle:
            random.shuffle(indices)
        self._send_py_dataset()  
        while True:
            with closing(self.executor_fn(_SHARED_SEQUENCES)) as executor:
                for i in indices:
                    if self.stop_signal.is_set():
                        return

                    self.queue.put(
                        executor.apply_async(get_index, (self.uid, i)),
                        block=True,
                    )

                
                self._wait_queue()

                if self.stop_signal.is_set():
                    
                    return

            
            self.py_dataset.on_epoch_end()
            self._send_py_dataset()  

    def get(self):
        
        while self.is_running():
            try:
                inputs = self.queue.get(block=True, timeout=5).get()
                if self.is_running():
                    self.queue.task_done()
                if inputs is not None:
                    yield inputs
            except queue.Empty:
                pass
            except Exception as e:
                self.stop()
                raise e


def init_pool_generator(gens, random_seed=None, id_queue=None):
    
    global _SHARED_SEQUENCES
    _SHARED_SEQUENCES = gens

    worker_proc = multiprocessing.current_process()

    
    
    worker_proc.name = f"Keras_worker_{worker_proc.name}"

    if random_seed is not None:
        np.random.seed(random_seed + worker_proc.ident)

    if id_queue is not None:
        
        id_queue.put(worker_proc.ident, block=True, timeout=0.1)