diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index 1384a2ea738..dccbc401432 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -145,6 +145,7 @@ py_library( "//tensorflow/lite/python:lite", "//tensorflow/python/compat", "//tensorflow/python/data", + "//tensorflow/python/distribute", "//tensorflow/python/distribute:estimator_training", "//tensorflow/python/eager:def_function", "//tensorflow/python/feature_column:feature_column_py", diff --git a/tensorflow/python/__init__.py b/tensorflow/python/__init__.py index 547043030b1..62a6e9f6d47 100644 --- a/tensorflow/python/__init__.py +++ b/tensorflow/python/__init__.py @@ -78,6 +78,7 @@ from tensorflow.python.ops import initializers_ns as initializers # Bring in subpackages. from tensorflow.python import data +from tensorflow.python import distribute from tensorflow.python import keras from tensorflow.python.feature_column import feature_column_lib as feature_column from tensorflow.python.layers import layers diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD index 2d9a1764db7..887c61cb8fd 100644 --- a/tensorflow/python/distribute/BUILD +++ b/tensorflow/python/distribute/BUILD @@ -103,6 +103,18 @@ cuda_py_test( ], ) +py_library( + name = "distribute", + srcs = [ + "__init__.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":distribute_lib", + ":mirrored_strategy", + ], +) + py_library( name = "distribute_lib", srcs = [ diff --git a/tensorflow/python/distribute/__init__.py b/tensorflow/python/distribute/__init__.py new file mode 100644 index 00000000000..4ff912ae10d --- /dev/null +++ b/tensorflow/python/distribute/__init__.py @@ -0,0 +1,25 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Distribution Strategy library.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# pylint: disable=unused-import +from tensorflow.python.distribute import distribute_lib +from tensorflow.python.distribute import distribution_strategy_context +from tensorflow.python.distribute import mirrored_strategy +# pylint: enable=unused-import diff --git a/tensorflow/python/distribute/mirrored_strategy.py b/tensorflow/python/distribute/mirrored_strategy.py index 0775920a7d6..7589720724c 100644 --- a/tensorflow/python/distribute/mirrored_strategy.py +++ b/tensorflow/python/distribute/mirrored_strategy.py @@ -43,6 +43,7 @@ from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import variable_scope from tensorflow.python.training import coordinator from tensorflow.python.util import nest +from tensorflow.python.util.tf_export import tf_export # TODO(josh11b): Replace asserts in this file with if ...: raise ... @@ -422,6 +423,7 @@ def all_local_devices(num_gpus=None): ("/device:CPU:0",)) +@tf_export("distribute.MirroredStrategy") class MirroredStrategy(distribute_lib.DistributionStrategy): """Mirrors vars to distribute across multiple devices and machines. diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt new file mode 100644 index 00000000000..a613e2d3d1d --- /dev/null +++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt @@ -0,0 +1,138 @@ +path: "tensorflow.distribute.MirroredStrategy" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "between_graph" + mtype: "" + } + member { + name: "extended" + mtype: "" + } + member { + name: "num_replicas_in_sync" + mtype: "" + } + member { + name: "parameter_devices" + mtype: "" + } + member { + name: "require_static_shapes" + mtype: "" + } + member { + name: "should_checkpoint" + mtype: "" + } + member { + name: "should_init" + mtype: "" + } + member { + name: "should_save_summary" + mtype: "" + } + member { + name: "worker_devices" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\', \'devices\', \'cross_device_ops\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " + } + member_method { + name: "batch_reduce" + argspec: "args=[\'self\', \'aggregation\', \'value_destination_pairs\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "broadcast" + argspec: "args=[\'self\', \'tensor\', \'destinations\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "call_for_each_replica" + argspec: "args=[\'self\', \'fn\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "colocate_vars_with" + argspec: "args=[\'self\', \'colocate_with_variable\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "configure" + argspec: "args=[\'self\', \'session_config\', \'cluster_spec\', \'task_type\', \'task_id\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], " + } + member_method { + name: "distribute_dataset" + argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "experimental_finalize" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "experimental_initialize" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "finalize" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "group" + argspec: "args=[\'self\', \'value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "initialize" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "make_dataset_iterator" + argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "make_input_fn_iterator" + argspec: "args=[\'self\', \'input_fn\', \'replication_mode\'], varargs=None, keywords=None, defaults=[\'InputReplicationMode.PER_WORKER\'], " + } + member_method { + name: "non_slot_devices" + argspec: "args=[\'self\', \'var_list\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "read_var" + argspec: "args=[\'self\', \'v\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "reduce" + argspec: "args=[\'self\', \'reduce_op\', \'value\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "run_steps_on_dataset" + argspec: "args=[\'self\', \'fn\', \'iterator\', \'iterations\', \'initial_loop_values\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], " + } + member_method { + name: "scope" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "unwrap" + argspec: "args=[\'self\', \'value\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "update" + argspec: "args=[\'self\', \'var\', \'fn\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "update_config_proto" + argspec: "args=[\'self\', \'config_proto\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "update_non_slot" + argspec: "args=[\'self\', \'colocate_with\', \'fn\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "value_container" + argspec: "args=[\'self\', \'value\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt index 4d833b54ba0..b0dd73ca1d4 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt @@ -8,6 +8,10 @@ tf_module { name: "InputReplicationMode" mtype: "" } + member { + name: "MirroredStrategy" + mtype: "" + } member { name: "ReduceOp" mtype: "" diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt new file mode 100644 index 00000000000..a613e2d3d1d --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt @@ -0,0 +1,138 @@ +path: "tensorflow.distribute.MirroredStrategy" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "between_graph" + mtype: "" + } + member { + name: "extended" + mtype: "" + } + member { + name: "num_replicas_in_sync" + mtype: "" + } + member { + name: "parameter_devices" + mtype: "" + } + member { + name: "require_static_shapes" + mtype: "" + } + member { + name: "should_checkpoint" + mtype: "" + } + member { + name: "should_init" + mtype: "" + } + member { + name: "should_save_summary" + mtype: "" + } + member { + name: "worker_devices" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\', \'devices\', \'cross_device_ops\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " + } + member_method { + name: "batch_reduce" + argspec: "args=[\'self\', \'aggregation\', \'value_destination_pairs\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "broadcast" + argspec: "args=[\'self\', \'tensor\', \'destinations\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "call_for_each_replica" + argspec: "args=[\'self\', \'fn\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "colocate_vars_with" + argspec: "args=[\'self\', \'colocate_with_variable\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "configure" + argspec: "args=[\'self\', \'session_config\', \'cluster_spec\', \'task_type\', \'task_id\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], " + } + member_method { + name: "distribute_dataset" + argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "experimental_finalize" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "experimental_initialize" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "finalize" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "group" + argspec: "args=[\'self\', \'value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "initialize" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "make_dataset_iterator" + argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "make_input_fn_iterator" + argspec: "args=[\'self\', \'input_fn\', \'replication_mode\'], varargs=None, keywords=None, defaults=[\'InputReplicationMode.PER_WORKER\'], " + } + member_method { + name: "non_slot_devices" + argspec: "args=[\'self\', \'var_list\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "read_var" + argspec: "args=[\'self\', \'v\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "reduce" + argspec: "args=[\'self\', \'reduce_op\', \'value\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "run_steps_on_dataset" + argspec: "args=[\'self\', \'fn\', \'iterator\', \'iterations\', \'initial_loop_values\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], " + } + member_method { + name: "scope" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "unwrap" + argspec: "args=[\'self\', \'value\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "update" + argspec: "args=[\'self\', \'var\', \'fn\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "update_config_proto" + argspec: "args=[\'self\', \'config_proto\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "update_non_slot" + argspec: "args=[\'self\', \'colocate_with\', \'fn\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "value_container" + argspec: "args=[\'self\', \'value\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt index 4d833b54ba0..b0dd73ca1d4 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt @@ -8,6 +8,10 @@ tf_module { name: "InputReplicationMode" mtype: "" } + member { + name: "MirroredStrategy" + mtype: "" + } member { name: "ReduceOp" mtype: ""