def pipe(max_batch_size, input_data, device):
pipe = Pipeline(batch_size=max_batch_size, num_threads=4, device_id=0)
data = fn.external_source(source=input_data,
cycle=False,
device=device)
processed = fn.lookup_table(data, keys=[1, 3], values=[10, 50])
pipe.set_outputs(processed)
return pipe
def test_lookup_table_cpu():
pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=None)
test_data_shape = [100]
def get_data():
out = [
np.random.randint(0, 5, size=test_data_shape, dtype=np.uint8)
for _ in range(batch_size)
]
return out
data = fn.external_source(source=get_data)
processed = fn.lookup_table(data, keys=[1, 3], values=[10, 50])
pipe.set_outputs(processed)
pipe.build()
for _ in range(3):
pipe.run()