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

Add colocation back for optimizer v2.

Browse Source 
PiperOrigin-RevId: 259403870
This commit is contained in:
Zhenyu Tan 2019-07-22 14:14:12 -07:00 committed by TensorFlower Gardener
parent e547d262a5
commit f8c912d2b7
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
tensorflow/python/keras/optimizer_v2/optimizer_v2.py
View File

@ -474,9 +474,12 @@ class OptimizerV2(trackable.Trackable):
update_ops = [] update_ops = []
with backend.name_scope(name or self._name): with backend.name_scope(name or self._name):
for grad, var in grads_and_vars: for grad, var in grads_and_vars:
scope_name = ("" if ops.executing_eagerly_outside_functions() else scope_name = ("update" if ops.executing_eagerly_outside_functions() else
"_" + var.op.name) "update_" + var.op.name)
with backend.name_scope("update" + scope_name): # Colocate the update with variables to avoid unnecessary communication
# delays. See b/136304694.
with backend.name_scope(
scope_name), distribution.extended.colocate_vars_with(var):
update_ops.extend( update_ops.extend(
distribution.extended.update( distribution.extended.update(
var, apply_grad_to_update_var, args=(grad,), group=False)) var, apply_grad_to_update_var, args=(grad,), group=False))