[tf.contrib.data] Avoid creating threads in the ParallelMap dataset.
This existing `ParallelMap(f, num_threads, output_buffer_size)` dataset has an iterator that uses `num_threads` threads to invoke `f`in parallel, and prefetch elements into a bounded buffer of max length `output_buffer_size`. This change modifies the implementation in two ways: * It replaces the `num_threads` worker threads with an equivalent number of concurrent async `Run()` calls on the underlying `FunctionLibraryRuntime`. * It factors out the prefetching behavior from ParallelMap, into a separate Prefetch dataset, as suggested by jsimsa@. To reflect the change in implementation, the `num_threads` input to `ParallelMapDatasetOp` is renamed to the more abstract `num_parallel_calls`. The net result is a reduction in the number of additional threads per parallel map from `num_threads` to 1 (if `output_buffer_size` is specified) or 0 (if `output_buffer_size` is not specified). Since these background threads offloaded work to the existing inter-op threadpool, there is no effective reduction in parallelism, and this approach should enable us to scale to larger degrees of parallelism because the system will have fewer threads to manage. PiperOrigin-RevId: 166154782
Loading
Please sign in to comment