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

Make sure that Adam colocates ops with a consistent variable across workers.

Browse Source 
PiperOrigin-RevId: 158022292
This commit is contained in:
Martin Wicke 2017-06-05 09:27:01 -07:00 committed by TensorFlower Gardener
parent 69ba4d3d49
commit 504a307b74
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tensorflow/python/training/adam.py
View File

@ -113,10 +113,14 @@ class AdamOptimizer(optimizer.Optimizer):
def _create_slots(self, var_list):
# Create the beta1 and beta2 accumulators on the same device as the first
# variable.
# variable. Sort the var_list to make sure this device is consistent across
# workers (these need to go on the same PS, otherwise some updates are
# silently ignored).
first_var = min(var_list, key=lambda x: x.name)
if (self._beta1_power is None or
self._beta1_power.graph is not var_list[0].graph):
with ops.colocate_with(var_list[0]):
self._beta1_power.graph is not first_var.graph):
with ops.colocate_with(first_var):
self._beta1_power = variable_scope.variable(self._beta1,
name="beta1_power",
trainable=False)