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