Function fast_large_random_batches

dali/test/python/operator_2/test_reduce.py:533–565

(rank, batch_size, num_batches, lo=0, hi=1)

def fast_large_random_batches(rank, batch_size, num_batches, lo=0, hi=1):
    max_vol = 10000000
    max_extent = min(65536, int(np.floor(max_vol ** (1 / rank))))

    # generate a maximum size buffer pre-filled with random numbers
    global _random_buf
    global _random_lo
    global _random_hi
    should_generate = (
        _random_buf is None
        or _random_buf.size < max_extent**rank
        or _random_lo != lo
        or _random_hi != hi
    )
    if should_generate:
        _random_lo = lo
        _random_hi = hi
        _random_buf = np.random.uniform(low=lo, high=hi, size=max_vol).astype(np.float32)

    data = []
    for _ in range(num_batches):
        batch = []
        for _ in range(batch_size):
            size = np.random.randint(1, max_extent, size=rank)
            vol = np.prod(size)
            # now that we know the actual volume of the sample, we can pick a random
            # location in the pre-filled buffer
            offset = np.random.randint(0, (_random_buf.size - vol) + 1)
            # take a slice and reshape it to the desired shape - these are constant time operations
            sample = _random_buf[offset : offset + vol].reshape(size)
            batch.append(sample)
        data.append(batch)
    return data


@nottest