PyDataset: choosing n of workers?

[santi-mir]

I'm thinking of using PyDataset but trying to understand first from first principles (before just testing) whether it makes sense to use a certain number of workers.

Here is an example from the docs:

// skip imports for brevity

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)

I'd guess that in this case the main reason to use multithreading and several workers would be to get a fast resizing but otherwise the operation will mostly involve I/O operations and not CPU processing.

So data-loading when a bit of processing is done to the data would be a good use case for multi threading & workers ? And what about a reasonable number of workers? Is there a way to know the maximum available?

[Kxrma47]

There is no universally correct worker count; it is the smallest count that keeps the accelerator/training step fed without creating CPU, memory, or storage contention.

In this example, image reading is I/O-heavy, while decoding and resizing may use CPU in native libraries. Threads can work well when those libraries release the GIL; multiprocessing is more useful for genuinely Python-bound CPU work, but has higher startup, serialization, and memory costs.

I would benchmark workers=1, 2, 4, 8 with the same epoch and watch training-step idle time, CPU utilization, storage throughput, and RAM. Stop increasing workers when examples/second no longer improves. os.cpu_count() is an upper bound on logical CPUs, not a recommended value: a network disk may saturate at two workers, while local NVMe plus expensive augmentation may benefit from more.

max_queue_size matters too. A queue of a few prefetched batches can hide occasional loader latency; making it huge mainly consumes memory.

If you enable multiprocessing, keep the dataset picklable, avoid sharing open file handles, and use the normal if __name__ == "__main__": guard on platforms that spawn processes. Also ensure __getitem__ is deterministic by batch index, because workers may request batches out of order even though Keras restores their logical order.