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")