def _test_external_source_callback_torch_stream(src_device, gen_device):
with torch.cuda.stream(torch.cuda.Stream()):
for attempt in range(10):
t0 = torch.tensor([attempt * 100 + 1.5],
dtype=torch.float32,
device=gen_device)
increment = torch.tensor([10],
dtype=torch.float32,
device=gen_device)
pipe = Pipeline(1, 3, 0)
def gen_batch():
nonlocal t0
t0 += increment
return [to_dlpack(t0)]
pipe.set_outputs(
fn.external_source(source=gen_batch,
device=src_device,
cuda_stream=torch.cuda.current_stream()))
pipe.build()
for i in range(10):
check_output(pipe.run(), [
np.array([attempt * 100 + (i + 1) * 10 + 1.5],
dtype=np.float32)
])
def test_external_source_collection_cycling_raise():
pipe = Pipeline(1, 3, 0, prefetch_queue_depth=1)
batches = [[make_array([1.5, 2.5], dtype=datapy.float32)],
[make_array([-1, 3.5, 4.5], dtype=datapy.float32)]]
def batch_gen():
for b in batches:
yield b
pipe.set_outputs(fn.external_source(batches, cycle="raise"),
fn.external_source(batch_gen, cycle="raise"))
pipe.build()
# epochs are cycles over the source iterable
for _ in range(3):
for batch in batches:
pipe_out = pipe.run()
batch = asnumpy(batch)
batch = batch, batch
check_output(pipe_out, batch)
with assert_raises(StopIteration):
pipe.run()
pipe.reset()
def test_external_source_with_iter():
pipe = Pipeline(1, 3, 0)
pipe.set_outputs(fn.external_source(lambda i: [np.array([i + 1.5], dtype=np.float32)]))
pipe.build()
for i in range(10):
check_output(pipe.run(), [np.array([i + 1.5], dtype=np.float32)])
def test_external_source_with_iter():
for attempt in range(10):
pipe = Pipeline(1, 3, 0)
pipe.set_outputs(fn.external_source(lambda i: [make_array([attempt * 100 + i * 10 + 1.5], dtype=datapy.float32)]))
pipe.build()
for i in range(10):
check_output(pipe.run(), [np.array([attempt * 100 + i * 10 + 1.5], dtype=np.float32)])
def run_and_check(pipe, ref_iterable):
iter_ref = iter(ref_iterable)
i = 0
while True:
try:
pipe_out = pipe.run()
check_output(pipe_out, next(iter_ref))
i += 1
except StopIteration:
break
assert (i == len(ref_iterable))
def test_external_source_generator():
pipe = Pipeline(1, 3, 0)
def gen():
for i in range(5):
yield [make_array([i + 1.5], dtype=datapy.float32)]
pipe.set_outputs(fn.external_source(gen()))
pipe.build()
for i in range(5):
check_output(pipe.run(), [np.array([i + 1.5], dtype=np.float32)])
def run_and_check(pipe, ref_iterable):
iter_ref = iter(ref_iterable)
i = 0
while True:
try:
pipe_out = pipe.run()
data = next(iter_ref)
data = asnumpy(data, iter_ref.device )
check_output(pipe_out, data)
i += 1
except StopIteration:
break
assert(i == len(ref_iterable))
def test_external_source_gen_function_cycle():
pipe = Pipeline(1, 3, 0)
def gen():
for i in range(5):
yield [make_array([i + 1.5], dtype=datapy.float32)]
pipe.set_outputs(fn.external_source(gen, cycle=True))
pipe.build()
for _ in range(3):
for i in range(5):
check_output(pipe.run(), [np.array([i + 1.5], dtype=np.float32)])
def test_external_source_collection_cycling():
pipe = Pipeline(1, 3, 0)
batches = [[np.array([1.5, 2.5], dtype=np.float32)],
[np.array([-1, 3.5, 4.5], dtype=np.float32)]]
pipe.set_outputs(fn.external_source(batches, cycle=True))
pipe.build()
# epochs are cycles over the source iterable
for epoch in range(3):
for batch in batches:
check_output(pipe.run(), batch)
def test_external_source_with_sample_info():
batch_size = 7
for attempt in range(10):
pipe = Pipeline(batch_size, 3, 0)
def src(si):
assert(si.idx_in_epoch == batch_size * si.iteration + si.idx_in_batch)
return make_array([attempt * 100 + si.iteration * 10 + si.idx_in_batch + 1.5], dtype=datapy.float32)
pipe.set_outputs(fn.external_source(src, batch=False))
pipe.build()
for i in range(10):
batch = [np.array([attempt * 100 + i * 10 + s + 1.5], dtype=np.float32) for s in range(batch_size)]
check_output(pipe.run(), batch)
def test_external_source_with_iter_cupy_stream():
with cp.cuda.Stream(non_blocking=True):
for attempt in range(10):
pipe = Pipeline(1, 3, 0)
def get_data(i):
return [
cp.array([attempt * 100 + i * 10 + 1.5], dtype=cp.float32)
]
pipe.set_outputs(fn.external_source(get_data))
pipe.build()
for i in range(10):
check_output(pipe.run(), [
np.array([attempt * 100 + i * 10 + 1.5], dtype=np.float32)
])
def test_external_source_callback_torch_stream():
with torch.cuda.stream(torch.cuda.Stream()):
for attempt in range(10):
t0 = torch.tensor([attempt * 100 + 1.5],
dtype=torch.float32).cuda()
increment = torch.tensor([10], dtype=torch.float32).cuda()
pipe = Pipeline(1, 3, 0)
def gen_batch():
nonlocal t0
t0 += increment
return [t0]
pipe.set_outputs(fn.external_source(gen_batch))
pipe.build()
for i in range(10):
check_output(pipe.run(), [
np.array([attempt * 100 + (i + 1) * 10 + 1.5],
dtype=np.float32)
])
