From d5d39154dbe8e6f7377ef10205039eda83a11fec Mon Sep 17 00:00:00 2001 From: Yuefeng Zhou Date: Wed, 20 Feb 2019 16:46:01 -0800 Subject: [PATCH] Expose MultiWorkerMirroredStrategy and ParameterServerStrategy PiperOrigin-RevId: 234895866 --- tensorflow/python/distribute/BUILD | 1 + tensorflow/python/distribute/__init__.py | 2 + .../collective_all_reduce_strategy.py | 2 + .../python/distribute/experimental/BUILD | 19 ++++++ .../distribute/experimental/__init__.py | 24 +++++++ .../distribute/parameter_server_strategy.py | 5 +- .../tools/api/generator/api_init_files.bzl | 1 + .../tools/api/generator/api_init_files_v1.bzl | 1 + ...ntal.-multi-worker-mirrored-strategy.pbtxt | 66 +++++++++++++++++++ ...erimental.-parameter-server-strategy.pbtxt | 66 +++++++++++++++++++ .../tensorflow.distribute.experimental.pbtxt | 11 ++++ .../api/golden/v1/tensorflow.distribute.pbtxt | 4 ++ ...ntal.-multi-worker-mirrored-strategy.pbtxt | 66 +++++++++++++++++++ ...erimental.-parameter-server-strategy.pbtxt | 66 +++++++++++++++++++ .../tensorflow.distribute.experimental.pbtxt | 11 ++++ .../api/golden/v2/tensorflow.distribute.pbtxt | 4 ++ 16 files changed, 348 insertions(+), 1 deletion(-) create mode 100644 tensorflow/python/distribute/experimental/BUILD create mode 100644 tensorflow/python/distribute/experimental/__init__.py create mode 100644 tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt create mode 100644 tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt create mode 100644 tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.pbtxt create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.pbtxt diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD index 3df97e14d19..feeae8d9f41 100644 --- a/tensorflow/python/distribute/BUILD +++ b/tensorflow/python/distribute/BUILD @@ -115,6 +115,7 @@ py_library( ":distribute_lib", ":mirrored_strategy", ":one_device_strategy", + "//tensorflow/python/distribute/experimental", ], ) diff --git a/tensorflow/python/distribute/__init__.py b/tensorflow/python/distribute/__init__.py index 87f746f9029..f9d0a95ea58 100644 --- a/tensorflow/python/distribute/__init__.py +++ b/tensorflow/python/distribute/__init__.py @@ -25,4 +25,6 @@ from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import distribution_strategy_context from tensorflow.python.distribute import mirrored_strategy from tensorflow.python.distribute import one_device_strategy +from tensorflow.python.distribute.experimental import collective_all_reduce_strategy +from tensorflow.python.distribute.experimental import parameter_server_strategy # pylint: enable=unused-import diff --git a/tensorflow/python/distribute/collective_all_reduce_strategy.py b/tensorflow/python/distribute/collective_all_reduce_strategy.py index 25bffe4015e..a9d092784d9 100644 --- a/tensorflow/python/distribute/collective_all_reduce_strategy.py +++ b/tensorflow/python/distribute/collective_all_reduce_strategy.py @@ -39,9 +39,11 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import collective_ops from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util.tf_export import tf_export # TODO(yuefengz): support in-graph replication. +@tf_export("distribute.experimental.MultiWorkerMirroredStrategy") class CollectiveAllReduceStrategy(distribute_lib.DistributionStrategy): """Distribution strategy that uses collective ops for all-reduce. diff --git a/tensorflow/python/distribute/experimental/BUILD b/tensorflow/python/distribute/experimental/BUILD new file mode 100644 index 00000000000..18893290f31 --- /dev/null +++ b/tensorflow/python/distribute/experimental/BUILD @@ -0,0 +1,19 @@ +package( + default_visibility = ["//tensorflow:internal"], +) + +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + +py_library( + name = "experimental", + srcs = [ + "__init__.py", + ], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python/distribute:collective_all_reduce_strategy", + "//tensorflow/python/distribute:parameter_server_strategy", + ], +) diff --git a/tensorflow/python/distribute/experimental/__init__.py b/tensorflow/python/distribute/experimental/__init__.py new file mode 100644 index 00000000000..f74cac8b5ff --- /dev/null +++ b/tensorflow/python/distribute/experimental/__init__.py @@ -0,0 +1,24 @@ +# 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. +# ============================================================================== +"""Experimental 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 collective_all_reduce_strategy +from tensorflow.python.distribute import parameter_server_strategy +# pylint: enable=unused-import diff --git a/tensorflow/python/distribute/parameter_server_strategy.py b/tensorflow/python/distribute/parameter_server_strategy.py index 7490eec4cf0..535327b9fe2 100644 --- a/tensorflow/python/distribute/parameter_server_strategy.py +++ b/tensorflow/python/distribute/parameter_server_strategy.py @@ -20,11 +20,12 @@ from __future__ import print_function import copy -from tensorflow.contrib.distribute.python import mirrored_strategy + from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib from tensorflow.python.distribute import device_util from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import input_lib +from tensorflow.python.distribute import mirrored_strategy from tensorflow.python.distribute import multi_worker_util from tensorflow.python.distribute import numpy_dataset from tensorflow.python.distribute import values @@ -39,12 +40,14 @@ from tensorflow.python.ops import variable_scope as vs from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import device_setter from tensorflow.python.util import nest +from tensorflow.python.util.tf_export import tf_export _LOCAL_CPU = "/device:CPU:0" _LOCAL_GPU_0 = "/device:GPU:0" # TODO(yuefengz): maybe cache variables on local CPU. +@tf_export("distribute.experimental.ParameterServerStrategy") class ParameterServerStrategy(distribute_lib.DistributionStrategy): """A parameter server DistributionStrategy. diff --git a/tensorflow/python/tools/api/generator/api_init_files.bzl b/tensorflow/python/tools/api/generator/api_init_files.bzl index 19af602cd23..9654a61a428 100644 --- a/tensorflow/python/tools/api/generator/api_init_files.bzl +++ b/tensorflow/python/tools/api/generator/api_init_files.bzl @@ -15,6 +15,7 @@ TENSORFLOW_API_INIT_FILES = [ "debugging/__init__.py", "distribute/__init__.py", "distribute/cluster_resolver/__init__.py", + "distribute/experimental/__init__.py", "dtypes/__init__.py", "errors/__init__.py", "experimental/__init__.py", diff --git a/tensorflow/python/tools/api/generator/api_init_files_v1.bzl b/tensorflow/python/tools/api/generator/api_init_files_v1.bzl index 01bf12ebe70..1179454d8b8 100644 --- a/tensorflow/python/tools/api/generator/api_init_files_v1.bzl +++ b/tensorflow/python/tools/api/generator/api_init_files_v1.bzl @@ -16,6 +16,7 @@ TENSORFLOW_API_INIT_FILES_V1 = [ "debugging/__init__.py", "distribute/__init__.py", "distribute/cluster_resolver/__init__.py", + "distribute/experimental/__init__.py", "distributions/__init__.py", "dtypes/__init__.py", "errors/__init__.py", diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt new file mode 100644 index 00000000000..619c3744d6c --- /dev/null +++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt @@ -0,0 +1,66 @@ +path: "tensorflow.distribute.experimental.MultiWorkerMirroredStrategy" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "extended" + mtype: "" + } + member { + name: "num_replicas_in_sync" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "broadcast" + argspec: "args=[\'self\', \'tensor\'], varargs=None, keywords=None, 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: "experimental_make_numpy_iterator" + argspec: "args=[\'self\', \'numpy_input\', \'batch_size\', \'num_epochs\', \'shuffle\', \'session\'], varargs=None, keywords=None, defaults=[\'1\', \'1024\', \'None\'], " + } + member_method { + name: "experimental_run" + argspec: "args=[\'self\', \'fn\', \'input_iterator\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "group" + argspec: "args=[\'self\', \'value\', \'name\'], 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: "reduce" + argspec: "args=[\'self\', \'reduce_op\', \'value\'], varargs=None, keywords=None, defaults=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_config_proto" + argspec: "args=[\'self\', \'config_proto\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt new file mode 100644 index 00000000000..c4cfa2f4ed2 --- /dev/null +++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt @@ -0,0 +1,66 @@ +path: "tensorflow.distribute.experimental.ParameterServerStrategy" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "extended" + mtype: "" + } + member { + name: "num_replicas_in_sync" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "broadcast" + argspec: "args=[\'self\', \'tensor\'], varargs=None, keywords=None, 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: "experimental_make_numpy_iterator" + argspec: "args=[\'self\', \'numpy_input\', \'batch_size\', \'num_epochs\', \'shuffle\', \'session\'], varargs=None, keywords=None, defaults=[\'1\', \'1024\', \'None\'], " + } + member_method { + name: "experimental_run" + argspec: "args=[\'self\', \'fn\', \'input_iterator\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "group" + argspec: "args=[\'self\', \'value\', \'name\'], 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: "reduce" + argspec: "args=[\'self\', \'reduce_op\', \'value\'], varargs=None, keywords=None, defaults=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_config_proto" + argspec: "args=[\'self\', \'config_proto\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.pbtxt new file mode 100644 index 00000000000..8859fbd38c4 --- /dev/null +++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.pbtxt @@ -0,0 +1,11 @@ +path: "tensorflow.distribute.experimental" +tf_module { + member { + name: "MultiWorkerMirroredStrategy" + mtype: "" + } + member { + name: "ParameterServerStrategy" + mtype: "" + } +} diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt index 5ea6b5b3489..430e8116606 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt @@ -56,6 +56,10 @@ tf_module { name: "cluster_resolver" mtype: "" } + member { + name: "experimental" + mtype: "" + } member_method { name: "get_loss_reduction" argspec: "args=[], varargs=None, keywords=None, defaults=None" diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt new file mode 100644 index 00000000000..619c3744d6c --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt @@ -0,0 +1,66 @@ +path: "tensorflow.distribute.experimental.MultiWorkerMirroredStrategy" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "extended" + mtype: "" + } + member { + name: "num_replicas_in_sync" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "broadcast" + argspec: "args=[\'self\', \'tensor\'], varargs=None, keywords=None, 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: "experimental_make_numpy_iterator" + argspec: "args=[\'self\', \'numpy_input\', \'batch_size\', \'num_epochs\', \'shuffle\', \'session\'], varargs=None, keywords=None, defaults=[\'1\', \'1024\', \'None\'], " + } + member_method { + name: "experimental_run" + argspec: "args=[\'self\', \'fn\', \'input_iterator\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "group" + argspec: "args=[\'self\', \'value\', \'name\'], 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: "reduce" + argspec: "args=[\'self\', \'reduce_op\', \'value\'], varargs=None, keywords=None, defaults=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_config_proto" + argspec: "args=[\'self\', \'config_proto\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt new file mode 100644 index 00000000000..c4cfa2f4ed2 --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt @@ -0,0 +1,66 @@ +path: "tensorflow.distribute.experimental.ParameterServerStrategy" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "extended" + mtype: "" + } + member { + name: "num_replicas_in_sync" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "broadcast" + argspec: "args=[\'self\', \'tensor\'], varargs=None, keywords=None, 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: "experimental_make_numpy_iterator" + argspec: "args=[\'self\', \'numpy_input\', \'batch_size\', \'num_epochs\', \'shuffle\', \'session\'], varargs=None, keywords=None, defaults=[\'1\', \'1024\', \'None\'], " + } + member_method { + name: "experimental_run" + argspec: "args=[\'self\', \'fn\', \'input_iterator\'], varargs=None, keywords=None, defaults=[\'None\'], " + } + member_method { + name: "group" + argspec: "args=[\'self\', \'value\', \'name\'], 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: "reduce" + argspec: "args=[\'self\', \'reduce_op\', \'value\'], varargs=None, keywords=None, defaults=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_config_proto" + argspec: "args=[\'self\', \'config_proto\'], varargs=None, keywords=None, defaults=None" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.pbtxt new file mode 100644 index 00000000000..8859fbd38c4 --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.pbtxt @@ -0,0 +1,11 @@ +path: "tensorflow.distribute.experimental" +tf_module { + member { + name: "MultiWorkerMirroredStrategy" + mtype: "" + } + member { + name: "ParameterServerStrategy" + mtype: "" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt index a4b78d4bb61..7339bee6cd8 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt @@ -56,6 +56,10 @@ tf_module { name: "cluster_resolver" mtype: "" } + member { + name: "experimental" + mtype: "" + } member_method { name: "get_replica_context" argspec: "args=[], varargs=None, keywords=None, defaults=None"