Home | History | Annotate | Download | only in ops 
      1 # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
      4 # you may not use this file except in compliance with the License.
      5 # You may obtain a copy of the License at
      6 #
      7 #     http://www.apache.org/licenses/LICENSE-2.0
      8 #
      9 # Unless required by applicable law or agreed to in writing, software
     10 # distributed under the License is distributed on an "AS IS" BASIS,
     11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 # See the License for the specific language governing permissions and
     13 # limitations under the License.
     14 # ==============================================================================
     15 """Experimental API for controlling optimizations in `tf.data` pipelines."""
     16 from __future__ import absolute_import
     17 from __future__ import division
     18 from __future__ import print_function
     19 
     20 
     21 from tensorflow.python.data.util import options
     22 from tensorflow.python.util.tf_export import tf_export
     23 
     24 
     25 @tf_export("data.experimental.OptimizationOptions")
     26 class OptimizationOptions(options.OptionsBase):
     27   """Represents options for dataset optimizations.
     28 
     29   You can set the optimization options of a dataset through the
     30   `experimental_optimization` property of `tf.data.Options`; the property is
     31   an instance of `tf.data.experimental.OptimizationOptions`.
     32 
     33   ```python
     34   options = tf.data.Options()
     35   options.experimental_optimization.map_vectorization = True
     36   options.experimental_optimization.apply_default_optimizations = False
     37   dataset = dataset.with_options(options)
     38   ```
     39   """
     40   apply_default_optimizations = options.create_option(
     41       name="apply_default_optimizations",
     42       ty=bool,
     43       docstring=
     44       "Whether to apply default static optimizations. If False, only static "
     45       "optimizations that have been explicitly enabled will be applied.")
     46 
     47   autotune = options.create_option(
     48       name="autotune",
     49       ty=bool,
     50       docstring=
     51       "Whether to automatically tune performance knobs. If None, defaults to "
     52       "True.")
     53 
     54   autotune_cpu_budget = options.create_option(
     55       name="autotune_cpu_budget",
     56       ty=int,
     57       docstring=
     58       "When autotuning is enabled (through `autotune`), determines the CPU "
     59       "budget to use. Values greater than the number of schedulable CPU cores "
     60       "are allowed but may result in CPU contention. If None, defaults to the "
     61       "number of schedulable CPU cores.")
     62 
     63   filter_fusion = options.create_option(
     64       name="filter_fusion",
     65       ty=bool,
     66       docstring=
     67       "Whether to fuse filter transformations. If None, defaults to False.")
     68 
     69   hoist_random_uniform = options.create_option(
     70       name="hoist_random_uniform",
     71       ty=bool,
     72       docstring=
     73       "Whether to hoist `tf.random_uniform()` ops out of map transformations. "
     74       "If None, defaults to False.")
     75 
     76   map_and_batch_fusion = options.create_option(
     77       name="map_and_batch_fusion",
     78       ty=bool,
     79       docstring=
     80       "Whether to fuse map and batch transformations. If None, defaults to "
     81       "True.")
     82 
     83   map_and_filter_fusion = options.create_option(
     84       name="map_and_filter_fusion",
     85       ty=bool,
     86       docstring=
     87       "Whether to fuse map and filter transformations. If None, defaults to "
     88       "False.")
     89 
     90   map_fusion = options.create_option(
     91       name="map_fusion",
     92       ty=bool,
     93       docstring="Whether to fuse map transformations. If None, defaults to "
     94       "False.")
     95 
     96   map_parallelization = options.create_option(
     97       name="map_parallelization",
     98       ty=bool,
     99       docstring=
    100       "Whether to parallelize stateless map transformations. If None, defaults "
    101       "to False.")
    102 
    103   map_vectorization = options.create_option(
    104       name="map_vectorization",
    105       ty=bool,
    106       docstring=
    107       "Whether to vectorize map transformations. If None, defaults to False.")
    108 
    109   noop_elimination = options.create_option(
    110       name="noop_elimination",
    111       ty=bool,
    112       docstring=
    113       "Whether to eliminate no-op transformations. If None, defaults to True.")
    114 
    115   shuffle_and_repeat_fusion = options.create_option(
    116       name="shuffle_and_repeat_fusion",
    117       ty=bool,
    118       docstring="Whether to fuse shuffle and repeat transformations. If None, "
    119       "defaults to True.")
    120 
    121   def _static_optimizations(self):
    122     """Produces the list of enabled static optimizations."""
    123     result = set()
    124     all_optimizations = [
    125         "filter_fusion",
    126         "hoist_random_uniform",
    127         "map_and_batch_fusion",
    128         "map_and_filter_fusion",
    129         "map_parallelization",
    130         "map_fusion",
    131         "map_vectorization",
    132         "noop_elimination",
    133         "shuffle_and_repeat_fusion",
    134     ]
    135     for optimization in all_optimizations:
    136       if getattr(self, optimization):
    137         result.add(optimization)
    138 
    139     if self.apply_default_optimizations is not False:
    140       # The following optimizations are turned on by default, unless the
    141       # user explicitly disables them.
    142       optimizations_to_disable = [
    143           "map_and_batch_fusion",
    144           "noop_elimination",
    145           "shuffle_and_repeat_fusion",
    146       ]
    147       for optimization in optimizations_to_disable:
    148         if getattr(self, optimization) is not False:
    149           result.add(optimization)
    150     return sorted(list(result))
    151