def pipe(max_batch_size, input_data, device): pipe = Pipeline(batch_size=max_batch_size, num_threads=4, device_id=0) perm = fn.batch_permutation(seed=420) data = fn.external_source(source=input_data, cycle=False, device=device) processed = fn.permute_batch(data, indices=perm) pipe.set_outputs(processed) return pipe
def test_batch_permute_cpu(): pipe = Pipeline(batch_size=batch_size, num_threads=3, device_id=None) data = fn.external_source(source=get_data, layout="HWC") perm = fn.batch_permutation(seed=420) processed = fn.permute_batch(data, indices=perm) pipe.set_outputs(processed) pipe.build() for _ in range(3): pipe.run()
def _test_permute_batch_out_of_range(device): batch_size = 10 pipe = Pipeline(batch_size, 4, 0) data = fn.external_source(source=lambda: gen_data(batch_size, np.int32), device=device, layout="abc") perm = fn.batch_permutation() pipe.set_outputs( data, fn.permute_batch(data, indices=[0, 1, 2, 3, 4, 5, 10, 7, 8, 9]), perm) pipe.build() pipe.run()
def _test_permute_batch_fixed(device): batch_size = 10 pipe = Pipeline(batch_size, 4, 0) data = fn.external_source(source=lambda: gen_data(batch_size, np.int16), device=device, layout="abc") idxs = [4, 8, 0, 6, 3, 5, 2, 9, 7, 1] pipe.set_outputs(data, fn.permute_batch(data, indices=idxs)) pipe.build() for i in range(10): orig, permuted = pipe.run() if isinstance(orig, dali.backend.TensorListGPU): orig = orig.as_cpu() ref = [orig.at(idx) for idx in idxs] check_batch(permuted, ref, len(ref), 0, 0, "abc")
def _test_permute_batch(device, type): batch_size = 10 pipe = Pipeline(batch_size, 4, 0) data = fn.external_source(source=lambda: gen_data(batch_size, type), device=device, layout="abc") perm = fn.batch_permutation() pipe.set_outputs(data, fn.permute_batch(data, indices=perm), perm) pipe.build() for i in range(10): orig, permuted, idxs = pipe.run() idxs = [int(idxs.at(i)) for i in range(batch_size)] if isinstance(orig, dali.backend.TensorListGPU): orig = orig.as_cpu() ref = [orig.at(idx) for idx in idxs] check_batch(permuted, ref, len(ref), 0, 0, "abc")