Class RandomDataIterator

dali/test/python/test_utils.py:413–440

class RandomDataIterator(object):
    def __init__(self, batch_size, shape=(10, 600, 800, 3), dtype=None, seed=0):
        import_numpy()
        # to avoid any numpy reference in the interface
        if dtype is None:
            dtype = np.uint8
        self.batch_size = batch_size
        self.test_data = []
        self.np_rng = np.random.default_rng(seed=seed)
        for _ in range(self.batch_size):
            if dtype == np.float32:
                self.test_data.append(
                    np.array(self.np_rng.random(shape) * (1.0), dtype=dtype) - 0.5
                )
            else:
                self.test_data.append(np.array(self.np_rng.random(shape) * 255, dtype=dtype))

    def __iter__(self):
        self.i = 0
        self.n = self.batch_size
        return self

    def __next__(self):
        batch = self.test_data
        self.i = (self.i + 1) % self.n
        return batch

    next = __next__


class RandomlyShapedDataIterator(object):