Function _test_sample_inflate

dali/test/python/operator_1/test_inflate.py:75–118

(op, batch_size, np_dtype, seed)

def _test_sample_inflate(op, batch_size, np_dtype, seed):
    epoch_size = 10 * batch_size
    rng = np.random.default_rng(seed=seed)
    permutation = rng.permutation(epoch_size)
    dtype = np_type_to_dali(np_dtype)

    def gen_iteration_sizes():
        num_yielded_samples = 0
        while num_yielded_samples < epoch_size:
            iteration_size = np.int32(np.floor(rng.uniform(1, batch_size + 1)))
            iteration_size = min(iteration_size, epoch_size - num_yielded_samples)
            yield iteration_size
            num_yielded_samples += iteration_size

    iteration_sizes = list(gen_iteration_sizes())
    assert sum(iteration_sizes) == epoch_size

    def source():
        num_yielded_samples = 0
        for iteration_size in iteration_sizes:
            sample_sizes = [permutation[num_yielded_samples + i] for i in range(iteration_size)]
            num_yielded_samples += iteration_size

            def sample(sample_size):
                start = (sample_size - 1) * sample_size // 2
                sample = np.arange(start, start + sample_size, dtype=np.int64).astype(np_dtype)
                return sample, sample_to_lz4(sample)

            samples, deflated = list(zip(*[sample(sample_size) for sample_size in sample_sizes]))
            yield list(samples), list(deflated), np.array(sample_sizes, dtype=np.int32)

    op, chck_ctx = _get_alias_and_ctx(op)

    @pipeline_def
    def pipeline():
        sample, deflated, shape = fn.external_source(source=source, batch=True, num_outputs=3)
        inflated = op(deflated.gpu(), shape=shape, dtype=dtype)
        return inflated, sample

    with chck_ctx:
        pipe = pipeline(batch_size=batch_size, num_threads=4, device_id=0)
        for iter_size in iteration_sizes:
            inflated, baseline = pipe.run()
            check_batch(inflated, baseline, iter_size)


@has_operator("decoders.inflate")