Function generate_data

dali/test/python/test_dali_variable_batch_size.py:72–118

Generates an epoch of data, that will be used for variable batch size verification. :param max_batch_size: Actual sizes of every batch in the epoch will be less or equal to max_batch_size :param n_iter: Number of iterations in the epoch :param sample_shap

(max_batch_size, n_iter, sample_shape, lo=0.0, hi=1.0, dtype=np.float32)

def generate_data(max_batch_size, n_iter, sample_shape, lo=0.0, hi=1.0, dtype=np.float32):
    """
    Generates an epoch of data, that will be used for variable batch size verification.

    :param max_batch_size: Actual sizes of every batch in the epoch will be less or equal
                           to max_batch_size
    :param n_iter: Number of iterations in the epoch
    :param sample_shape: If sample_shape is callable, shape of every sample will be determined by
                         calling sample_shape. In this case, every call to sample_shape has to
                         return a tuple of integers. If sample_shape is a tuple, this will be a
                         shape of every sample.
    :param lo: Begin of the random range
    :param hi: End of the random range
    :param dtype: Numpy data type
    :return: An epoch of data
    """
    batch_sizes = np.array([max_batch_size // 2, max_batch_size // 4, max_batch_size])

    if isinstance(sample_shape, tuple):

        def sample_shape_wrapper():
            return sample_shape

        size_fn = sample_shape_wrapper
    elif inspect.isfunction(sample_shape):
        size_fn = sample_shape
    else:
        raise RuntimeError(
            "`sample_shape` shall be either a tuple or a callable. "
            "Provide `(val,)` tuple for 1D shape"
        )

    if np.issubdtype(dtype, np.integer):
        return [
            np.random.randint(lo, hi, size=(bs,) + size_fn(), dtype=dtype) for bs in batch_sizes
        ]
    elif np.issubdtype(dtype, np.float32):
        ret = (np.random.random_sample(size=(bs,) + size_fn()) for bs in batch_sizes)
        ret = map(lambda batch: (hi - lo) * batch + lo, ret)
        ret = map(lambda batch: batch.astype(dtype), ret)
        return list(ret)
    elif np.issubdtype(dtype, bool):
        assert isinstance(lo, bool)
        assert isinstance(hi, bool)
        return [np.random.choice(a=[lo, hi], size=(bs,) + size_fn()) for bs in batch_sizes]
    else:
        raise RuntimeError(f"Invalid type argument: {dtype}")


def single_op_pipeline(