From 14a437741fbb29a335903f59680d9679b249b3cd Mon Sep 17 00:00:00 2001 From: Yibo Zhu Date: Fri, 28 Jun 2019 13:18:39 +0800 Subject: [PATCH] tensorflow: better handling name_scope --- byteps/tensorflow/__init__.py | 10 ++++++---- byteps/tensorflow/ops.py | 10 ++++++++-- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/byteps/tensorflow/__init__.py b/byteps/tensorflow/__init__.py index 58bf897a3..456bbacb3 100644 --- a/byteps/tensorflow/__init__.py +++ b/byteps/tensorflow/__init__.py @@ -56,23 +56,25 @@ def push_pull(tensor, scope='', average=True, device_dense='', device_sparse='', return new_tensor -def broadcast_global_variables(root_rank): +def broadcast_global_variables(root_rank, scope=''): """Broadcasts all global variables from root rank to all other processes. Arguments: root_rank: rank of the process from which global variables will be broadcasted to all other processes. + scope: the graph name scope """ - return broadcast_variables(tf.global_variables(), root_rank) + return broadcast_variables(tf.global_variables(), root_rank, scope) -def broadcast_variables(variables, root_rank): +def broadcast_variables(variables, root_rank, scope=''): """Broadcasts variables from root rank to all other processes. Arguments: variables: variables for broadcast root_rank: rank of the process from which global variables will be broadcasted to all other processes. + scope: the graph name scope """ - return tf.group(*[tf.assign(var, broadcast(var, root_rank)) + return tf.group(*[tf.assign(var, broadcast(var, root_rank, scope)) for var in variables]) diff --git a/byteps/tensorflow/ops.py b/byteps/tensorflow/ops.py index aceebf396..870fbaf76 100644 --- a/byteps/tensorflow/ops.py +++ b/byteps/tensorflow/ops.py @@ -79,6 +79,8 @@ def _push_pull(tensor, scope='', name=None): """ if name is None and not _executing_eagerly(): name = 'BytePSPushPull_%s' % _normalize_name(tensor.name) + if scope == '' and not _executing_eagerly(): + scope = tf.compat.v1.get_default_graph().get_name_scope() full_name = scope + name full_name = full_name.encode("ascii") TF_LIB_CTYPES.byteps_tensorflow_declare_tensor(ctypes.c_char_p(full_name)) @@ -97,7 +99,7 @@ def _push_pull_grad(op, grad): return _push_pull(grad) -def broadcast(tensor, root_rank, name=None, is_variable=True): +def broadcast(tensor, root_rank, scope='', name=None, is_variable=True): """An op which broadcasts the input tensor on root rank to the same input tensor on all other BytePS processes. The broadcast operation is keyed by the name of the op. The tensor type and @@ -110,7 +112,11 @@ def broadcast(tensor, root_rank, name=None, is_variable=True): # Broadcast is implemented as push + pull after zero-ing non-root tensors if name is None and not _executing_eagerly(): name = 'BytePSBroadcast_%s' % _normalize_name(tensor.name) - TF_LIB_CTYPES.byteps_tensorflow_declare_tensor(ctypes.c_char_p(name.encode("ascii"))) + if scope == '' and not _executing_eagerly(): + scope = tf.compat.v1.get_default_graph().get_name_scope() + full_name = scope + name + full_name = full_name.encode("ascii") + TF_LIB_CTYPES.byteps_tensorflow_declare_tensor(ctypes.c_char_p(full_name)) if is_variable and (root_rank != rank()): return C_LIB.byteps_push_pull(tensor.assign(tf.zeros_like(tensor)), name=name) else: