experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[tf.data] Give users the option to input RAM budget in autotuning.

Browse Source 
PiperOrigin-RevId: 329784678
Change-Id: I3ad64c99aae42906f1ed96c46dc0885b518978fb
This commit is contained in:
Jay Shi 2020-09-02 14:08:57 -07:00 committed by TensorFlower Gardener
parent 768c68ceae
commit 67dfaf3a87
9 changed files with 47 additions and 13 deletions

12
tensorflow/core/kernels/data/model_dataset_op.cc
View File

@ -50,7 +50,17 @@ class ModelDatasetOp : public UnaryDatasetOpKernel {
OP_REQUIRES(ctx, cpu_budget_ > 0,
errors::InvalidArgument("CPU budget must be positive but is ",
cpu_budget_, "."));
ram_budget_ = kRamBudgetShare * port::AvailableRam();
if (ctx->HasAttr("ram_budget")) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("ram_budget", &ram_budget_));
} else {
ram_budget_ = 0;
}
if (ram_budget_ == 0) {
ram_budget_ = kRamBudgetShare * port::AvailableRam();
}
OP_REQUIRES(ctx, ram_budget_ > 0,
errors::InvalidArgument("RAM budget must be positive but is ",
ram_budget_, "."));
}
void MakeDataset(OpKernelContext* ctx, DatasetBase* input,

1
tensorflow/core/ops/dataset_ops.cc
View File

@ -882,6 +882,7 @@ REGISTER_OP("ModelDataset")
.Output("handle: variant")
.Attr("algorithm: int = 0")
.Attr("cpu_budget: int = 0")
.Attr("ram_budget: int = 0")
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
.SetShapeFn(shape_inference::ScalarShape);

13
tensorflow/python/data/experimental/kernel_tests/optimize_dataset_test.py
View File

@ -429,23 +429,26 @@ class OptimizeDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
options = dataset_ops.Options()
# Check defaults
autotune, algorithm, cpu_budget = options._autotune_settings()
autotune, algorithm, cpu_budget, ram_budget = options._autotune_settings()
self.assertTrue(autotune)
self.assertEqual(algorithm,
optimization_options._AutotuneAlgorithm.HILL_CLIMB)
self.assertEqual(cpu_budget, 0)
self.assertEqual(ram_budget, 0)
@combinations.generate(test_base.default_test_combinations())
def testAutotuningBufferSizes(self):
def testAutotuningSettings(self):
options = dataset_ops.Options()
options.experimental_optimization.autotune_cpu_budget = 1000
options.experimental_optimization.autotune_ram_budget = 999999999
options.experimental_optimization.autotune_buffers = True
self.assertIn("inject_prefetch", options._graph_rewrites().enabled)
autotune, algorithm, cpu_budget = options._autotune_settings()
autotune, algorithm, cpu_budget, ram_budget = options._autotune_settings()
self.assertTrue(autotune)
self.assertEqual(algorithm,
optimization_options._AutotuneAlgorithm.GRADIENT_DESCENT)
self.assertEqual(cpu_budget, 0)
self.assertEqual(cpu_budget, 1000)
self.assertEqual(ram_budget, 999999999)
if __name__ == "__main__":
test.main()

13
tensorflow/python/data/experimental/ops/optimization_options.py
View File

@ -117,6 +117,14 @@ class OptimizationOptions(options.OptionsBase):
"are allowed but may result in CPU contention. If None, defaults to the "
"number of schedulable CPU cores.")
autotune_ram_budget = options.create_option(
name="autotune_ram_budget",
ty=int,
docstring=
"When autotuning is enabled (through `autotune`), determines the RAM "
"budget to use. Values greater than the available RAM in bytes may "
"result in OOM. If None, defaults to half of the available RAM in bytes.")
filter_fusion = options.create_option(
name="filter_fusion",
ty=bool,
@ -223,14 +231,17 @@ class OptimizationOptions(options.OptionsBase):
_AutotuneAlgorithm.GRADIENT_DESCENT
if self._autotune_buffers() else _AutotuneAlgorithm.HILL_CLIMB)
cpu_budget = 0 # Indicates that all CPU cores should be used by default.
ram_budget = 0 # Indicates that default value of RAM budget should be used.
# Set these options if they are explicitly set by the user.
if self.autotune is False: # pylint: disable=g-bool-id-comparison
autotune = False
if self.autotune_cpu_budget is not None:
cpu_budget = self.autotune_cpu_budget
if self.autotune_ram_budget is not None:
ram_budget = self.autotune_ram_budget
return autotune, algorithm, cpu_budget
return autotune, algorithm, cpu_budget, ram_budget
def _graph_rewrites(self):
"""Produces lists of enabled, disabled and default graph optimizations.

7
tensorflow/python/data/ops/dataset_ops.py
View File

@ -392,10 +392,10 @@ class DatasetV2(collections_abc.Iterable, tracking_base.Trackable,
graph_rewrites.default, graph_rewrite_configs)
# (3) Apply autotune options
autotune, algorithm, cpu_budget = options._autotune_settings() # pylint: disable=protected-access
autotune, algorithm, cpu_budget, ram_budget = options._autotune_settings() # pylint: disable=protected-access
if autotune:
dataset = _ModelDataset(dataset, algorithm, cpu_budget)
dataset = _ModelDataset(dataset, algorithm, cpu_budget, ram_budget)
# (4) Apply stats aggregator options
if options.experimental_stats and options.experimental_stats.aggregator: # pylint: disable=line-too-long
@ -4445,12 +4445,13 @@ class _OptionsDataset(UnaryUnchangedStructureDataset):
class _ModelDataset(UnaryUnchangedStructureDataset):
"""A `Dataset` that acts as an identity, and models performance."""
def __init__(self, input_dataset, algorithm, cpu_budget):
def __init__(self, input_dataset, algorithm, cpu_budget, ram_budget):
self._input_dataset = input_dataset
variant_tensor = gen_dataset_ops.model_dataset(
input_dataset._variant_tensor, # pylint: disable=protected-access
algorithm=algorithm.value,
cpu_budget=cpu_budget,
ram_budget=ram_budget,
**self._flat_structure)
super(_ModelDataset, self).__init__(input_dataset, variant_tensor)

4
tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optimization-options.pbtxt
View File

@ -19,6 +19,10 @@ tf_class {
name: "autotune_cpu_budget"
mtype: "<type \'property\'>"
}
member {
name: "autotune_ram_budget"
mtype: "<type \'property\'>"
}
member {
name: "filter_fusion"
mtype: "<type \'property\'>"

2
tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
View File

@ -2518,7 +2518,7 @@ tf_module {
}
member_method {
name: "ModelDataset"
argspec: "args=[\'input_dataset\', \'output_types\', \'output_shapes\', \'algorithm\', \'cpu_budget\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'0\', \'None\'], "
argspec: "args=[\'input_dataset\', \'output_types\', \'output_shapes\', \'algorithm\', \'cpu_budget\', \'ram_budget\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'0\', \'0\', \'None\'], "
}
member_method {
name: "Mul"

4
tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optimization-options.pbtxt
View File

@ -19,6 +19,10 @@ tf_class {
name: "autotune_cpu_budget"
mtype: "<type \'property\'>"
}
member {
name: "autotune_ram_budget"
mtype: "<type \'property\'>"
}
member {
name: "filter_fusion"
mtype: "<type \'property\'>"

4
tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
View File

@ -1883,7 +1883,7 @@ tf_module {
member_method {
name: "ImageProjectiveTransformV3"
argspec: "args=[\'images\', \'transforms\', \'output_shape\', \'fill_value\', \'interpolation\', \'fill_mode\', \'name\'], varargs=None, keywords=None, defaults=[\'CONSTANT\', \'None\'], "
}
}
member_method {
name: "ImageSummary"
argspec: "args=[\'tag\', \'tensor\', \'max_images\', \'bad_color\', \'name\'], varargs=None, keywords=None, defaults=[\'3\', \'dtype: DT_UINT8\\ntensor_shape {\\n dim {\\n size: 4\\n }\\n}\\nint_val: 255\\nint_val: 0\\nint_val: 0\\nint_val: 255\\n\', \'None\'], "
@ -2518,7 +2518,7 @@ tf_module {
}
member_method {
name: "ModelDataset"
argspec: "args=[\'input_dataset\', \'output_types\', \'output_shapes\', \'algorithm\', \'cpu_budget\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'0\', \'None\'], "
argspec: "args=[\'input_dataset\', \'output_types\', \'output_shapes\', \'algorithm\', \'cpu_budget\', \'ram_budget\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'0\', \'0\', \'None\'], "
}
member_method {
name: "Mul"