easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/compat/sync_replicas_optimizer.py

@@ -0,0 +1,528 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Synchronize replicas for training."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.core.framework import types_pb2
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import optimizer
+from tensorflow.python.training import queue_runner
+from tensorflow.python.training import session_manager
+from tensorflow.python.training import session_run_hook
+from tensorflow.python.util.tf_export import tf_export
+
+
+# Please note that the gradients from replicas are averaged instead of summed
+# (as in the old sync_replicas_optimizer) so you need to increase the learning
+# rate according to the number of replicas. This change is introduced to be
+# consistent with how gradients are aggregated (averaged) within a batch in a
+# replica.
+@tf_export('train.SyncReplicasOptimizer')
+class SyncReplicasOptimizer(optimizer.Optimizer):
+  """Class to synchronize, aggregate gradients and pass them to the optimizer.
+
+  In a typical asynchronous training environment, it's common to have some
+  stale gradients. For example, with a N-replica asynchronous training,
+  gradients will be applied to the variables N times independently. Depending
+  on each replica's training speed, some gradients might be calculated from
+  copies of the variable from several steps back (N-1 steps on average). This
+  optimizer avoids stale gradients by collecting gradients from all replicas,
+  averaging them, then applying them to the variables in one shot, after
+  which replicas can fetch the new variables and continue.
+
+  The following accumulators/queue are created:
+
+  * N `gradient accumulators`, one per variable to train. Gradients are pushed
+    to them and the chief worker will wait until enough gradients are collected
+    and then average them before applying to variables. The accumulator will
+    drop all stale gradients (more details in the accumulator op).
+  * 1 `token` queue where the optimizer pushes the new global_step value after
+    all variables are updated.
+
+  The following local variable is created:
+  * `sync_rep_local_step`, one per replica. Compared against the global_step in
+    each accumulator to check for staleness of the gradients.
+
+  The optimizer adds nodes to the graph to collect gradients and pause the
+  trainers until variables are updated.
+  For the Parameter Server job:
+
+  1. An accumulator is created for each variable, and each replica pushes the
+     gradients into the accumulators instead of directly applying them to the
+     variables.
+  2. Each accumulator averages once enough gradients (replicas_to_aggregate)
+     have been accumulated.
+  3. Apply the averaged gradients to the variables.
+  4. Only after all variables have been updated, increment the global step.
+  5. Only after step 4, pushes `global_step` in the `token_queue`, once for
+     each worker replica. The workers can now fetch the global step, use it to
+     update its local_step variable and start the next batch.
+
+  For the replicas:
+
+  1. Start a step: fetch variables and compute gradients.
+  2. Once the gradients have been computed, push them into gradient
+     accumulators. Each accumulator will check the staleness and drop the stale.
+  3. After pushing all the gradients, dequeue an updated value of global_step
+     from the token queue and record that step to its local_step variable. Note
+     that this is effectively a barrier.
+  4. Start the next batch.
+
+  ### Usage
+
+  ```python
+  # Create any optimizer to update the variables, say a simple SGD:
+  opt = GradientDescentOptimizer(learning_rate=0.1)
+
+  # Wrap the optimizer with sync_replicas_optimizer with 50 replicas: at each
+  # step the optimizer collects 50 gradients before applying to variables.
+  # Note that if you want to have 2 backup replicas, you can change
+  # total_num_replicas=52 and make sure this number matches how many physical
+  # replicas you started in your job.
+  opt = tf.train.SyncReplicasOptimizer(opt, replicas_to_aggregate=50,
+                                 total_num_replicas=50)
+
+  # Some models have startup_delays to help stabilize the model but when using
+  # sync_replicas training, set it to 0.
+
+  # Now you can call `minimize()` or `compute_gradients()` and
+  # `apply_gradients()` normally
+  training_op = opt.minimize(total_loss, global_step=self.global_step)
+
+
+  # You can create the hook which handles initialization and queues.
+  sync_replicas_hook = opt.make_session_run_hook(is_chief)
+  ```
+
+  In the training program, every worker will run the train_op as if not
+  synchronized.
+
+  ```python
+  with training.MonitoredTrainingSession(
+      master=workers[worker_id].target, is_chief=is_chief,
+      hooks=[sync_replicas_hook]) as mon_sess:
+    while not mon_sess.should_stop():
+      mon_sess.run(training_op)
+  ```
+
+  To use SyncReplicasOptimizer with an `Estimator`, you need to send
+  sync_replicas_hook while calling the fit.
+  ```python
+  my_estimator = DNNClassifier(..., optimizer=opt)
+  my_estimator.fit(..., hooks=[sync_replicas_hook])
+  ```
+  """
+
+  sync_que_id = -1
+
+  def __init__(self,
+               opt,
+               replicas_to_aggregate,
+               total_num_replicas=None,
+               variable_averages=None,
+               variables_to_average=None,
+               use_locking=False,
+               name='sync_replicas',
+               **extra_args):
+    """Construct a sync_replicas optimizer.
+
+    Args:
+      opt: The actual optimizer that will be used to compute and apply the
+        gradients. Must be one of the Optimizer classes.
+      replicas_to_aggregate: number of replicas to aggregate for each variable
+        update.
+      total_num_replicas: Total number of tasks/workers/replicas, could be
+        different from replicas_to_aggregate.
+        If total_num_replicas > replicas_to_aggregate: it is backup_replicas +
+        replicas_to_aggregate.
+        If total_num_replicas < replicas_to_aggregate: Replicas compute
+        multiple batches per update to variables.
+      variable_averages: Optional `ExponentialMovingAverage` object, used to
+        maintain moving averages for the variables passed in
+        `variables_to_average`.
+      variables_to_average: a list of variables that need to be averaged. Only
+        needed if variable_averages is passed in.
+      use_locking: If True use locks for update operation.
+      name: string. Optional name of the returned operation.
+    """
+    if total_num_replicas is None:
+      total_num_replicas = replicas_to_aggregate
+
+    super(SyncReplicasOptimizer, self).__init__(use_locking, name)
+    logging.info(
+        'SyncReplicasV2: replicas_to_aggregate=%s; total_num_replicas=%s',
+        replicas_to_aggregate, total_num_replicas)
+    self._opt = opt
+    self._replicas_to_aggregate = replicas_to_aggregate
+    self._gradients_applied = False
+    self._variable_averages = variable_averages
+    self._variables_to_average = variables_to_average
+    self._total_num_replicas = total_num_replicas
+    self._tokens_per_step = max(total_num_replicas, replicas_to_aggregate)
+    self._global_step = None
+    self._sync_token_queue = None
+    self._is_sync_que_closed = None
+    self._close_sync_que = None
+
+    # The synchronization op will be executed in a queue runner which should
+    # only be executed by one of the replicas (usually the chief).
+    self._chief_queue_runner = None
+
+    # Remember which accumulator is on which device to set the initial step in
+    # the accumulator to be global step. This list contains list of the
+    # following format: (accumulator, device).
+    self._accumulator_list = []
+
+  def compute_gradients(self, *args, **kwargs):
+    """Compute gradients of "loss" for the variables in "var_list".
+
+    This simply wraps the compute_gradients() from the real optimizer. The
+    gradients will be aggregated in the apply_gradients() so that user can
+    modify the gradients like clipping with per replica global norm if needed.
+    The global norm with aggregated gradients can be bad as one replica's huge
+    gradients can hurt the gradients from other replicas.
+
+    Args:
+      *args: Arguments for compute_gradients().
+      **kwargs: Keyword arguments for compute_gradients().
+
+    Returns:
+      A list of (gradient, variable) pairs.
+    """
+    return self._opt.compute_gradients(*args, **kwargs)
+
+  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
+    """Apply gradients to variables.
+
+    This contains most of the synchronization implementation and also wraps the
+    apply_gradients() from the real optimizer.
+
+    Args:
+      grads_and_vars: List of (gradient, variable) pairs as returned by
+        compute_gradients().
+      global_step: Optional Variable to increment by one after the
+        variables have been updated.
+      name: Optional name for the returned operation.  Default to the
+        name passed to the Optimizer constructor.
+
+    Returns:
+      train_op: The op to dequeue a token so the replicas can exit this batch
+      and start the next one. This is executed by each replica.
+
+    Raises:
+      ValueError: If the grads_and_vars is empty.
+      ValueError: If global step is not provided, the staleness cannot be
+        checked.
+    """
+    if not grads_and_vars:
+      raise ValueError('Must supply at least one variable')
+
+    if global_step is None:
+      raise ValueError('Global step is required to check staleness')
+
+    self._global_step = global_step
+    train_ops = []
+    aggregated_grad = []
+    var_list = []
+
+    # local_anchor op will be placed on this worker task by default.
+    local_anchor = control_flow_ops.no_op()
+    # Colocating local_step variable prevents it being placed on the PS.
+    with ops.colocate_with(local_anchor):
+      self._local_step = variable_scope.variable(
+          initial_value=0,
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          dtype=global_step.dtype.base_dtype,
+          name='sync_rep_local_step')
+
+    self.local_step_init_op = state_ops.assign(self._local_step, global_step)
+    chief_init_ops = [self.local_step_init_op]
+    self.ready_for_local_init_op = variables.report_uninitialized_variables(
+        variables.global_variables())
+
+    with ops.name_scope(None, self._name):
+      for grad, var in grads_and_vars:
+        var_list.append(var)
+        with ops.device(var.device):
+          # Dense gradients.
+          if grad is None:
+            aggregated_grad.append(None)  # pass-through.
+            continue
+          elif isinstance(grad, ops.Tensor):
+            grad_accum = data_flow_ops.ConditionalAccumulator(
+                grad.dtype,
+                shape=var.get_shape(),
+                shared_name=var.name + '/grad_accum')
+            train_ops.append(
+                grad_accum.apply_grad(grad, local_step=self._local_step))
+            aggregated_grad.append(
+                grad_accum.take_grad(self._replicas_to_aggregate))
+          else:
+            if not isinstance(grad, ops.IndexedSlices):
+              raise ValueError('Unknown grad type!')
+            grad_accum = data_flow_ops.SparseConditionalAccumulator(
+                grad.dtype, shape=(), shared_name=var.name + '/grad_accum')
+            train_ops.append(
+                grad_accum.apply_indexed_slices_grad(
+                    grad, local_step=self._local_step))
+            aggregated_grad.append(
+                grad_accum.take_indexed_slices_grad(
+                    self._replicas_to_aggregate))
+
+          self._accumulator_list.append((grad_accum, var.device))
+
+      aggregated_grads_and_vars = zip(aggregated_grad, var_list)
+
+      # sync_op will be assigned to the same device as the global step.
+      with ops.device(global_step.device), ops.name_scope(''):
+        update_op = self._opt.apply_gradients(aggregated_grads_and_vars,
+                                              global_step)
+
+      def _get_token_qname():
+        SyncReplicasOptimizer.sync_que_id += 1
+        if SyncReplicasOptimizer.sync_que_id == 0:
+          return 'sync_token_q'
+        else:
+          return 'sync_token_q_' + str(SyncReplicasOptimizer.sync_que_id)
+
+      # Create token queue.
+      token_qname = _get_token_qname()
+      logging.info('create sync_token_queue[%s]' % token_qname)
+      with ops.device(global_step.device), ops.name_scope(''):
+        sync_token_queue = (
+            data_flow_ops.FIFOQueue(
+                -1,
+                global_step.dtype.base_dtype,
+                shapes=(),
+                name=token_qname,
+                shared_name=token_qname))
+        self._sync_token_queue = sync_token_queue
+        self._is_sync_que_closed = sync_token_queue.is_closed()
+        self._close_sync_que = sync_token_queue.close(
+            cancel_pending_enqueues=True, name='close_sync_token_queue')
+
+        # dummy_queue is passed to the queue runner. Don't use the real queues
+        # because the queue runner doesn't automatically reopen it once it
+        # closed queues in PS devices.
+        dummy_queue = (
+            data_flow_ops.FIFOQueue(
+                1,
+                types_pb2.DT_INT32,
+                shapes=(),
+                name='dummy_queue',
+                shared_name='dummy_queue'))
+
+      with ops.device(global_step.device), ops.name_scope(''):
+        # Replicas have to wait until they can get a token from the token queue.
+        with ops.control_dependencies(train_ops):
+          token = sync_token_queue.dequeue()
+        train_op = state_ops.assign(self._local_step, token)
+
+        with ops.control_dependencies([update_op]):
+          # Sync_op needs to insert tokens to the token queue at the end of the
+          # step so the replicas can fetch them to start the next step.
+          tokens = array_ops.fill([self._tokens_per_step], global_step)
+          sync_op = sync_token_queue.enqueue_many((tokens,))
+
+        if self._variable_averages is not None:
+          with ops.control_dependencies([sync_op]), ops.name_scope(''):
+            sync_op = self._variable_averages.apply(self._variables_to_average)
+
+        self._chief_queue_runner = queue_runner.QueueRunner(
+            dummy_queue, [sync_op])
+        ops.add_to_collection(ops.GraphKeys.QUEUE_RUNNERS,
+                              self._chief_queue_runner)
+      for accum, dev in self._accumulator_list:
+        with ops.device(dev):
+          chief_init_ops.append(
+              accum.set_global_step(global_step, name='SetGlobalStep'))
+      self.chief_init_op = control_flow_ops.group(*(chief_init_ops))
+      self._gradients_applied = True
+      return train_op
+
+  def get_chief_queue_runner(self):
+    """Returns the QueueRunner for the chief to execute.
+
+    This includes the operations to synchronize replicas: aggregate gradients,
+    apply to variables, increment global step, insert tokens to token queue.
+
+    Note that this can only be called after calling apply_gradients() which
+    actually generates this queuerunner.
+
+    Returns:
+      A `QueueRunner` for chief to execute.
+
+    Raises:
+      ValueError: If this is called before apply_gradients().
+    """
+    if self._gradients_applied is False:
+      raise ValueError('Should be called after apply_gradients().')
+
+    return self._chief_queue_runner
+
+  def get_slot(self, *args, **kwargs):
+    """Return a slot named "name" created for "var" by the Optimizer.
+
+    This simply wraps the get_slot() from the actual optimizer.
+
+    Args:
+      *args: Arguments for get_slot().
+      **kwargs: Keyword arguments for get_slot().
+
+    Returns:
+      The `Variable` for the slot if it was created, `None` otherwise.
+    """
+    return self._opt.get_slot(*args, **kwargs)
+
+  def variables(self):
+    """Fetches a list of optimizer variables in the default graph.
+
+    This wraps `variables()` from the actual optimizer. It does not include
+    the `SyncReplicasOptimizer`'s local step.
+
+    Returns:
+      A list of variables.
+    """
+    return self._opt.variables()
+
+  def get_slot_names(self, *args, **kwargs):
+    """Return a list of the names of slots created by the `Optimizer`.
+
+    This simply wraps the get_slot_names() from the actual optimizer.
+
+    Args:
+      *args: Arguments for get_slot().
+      **kwargs: Keyword arguments for get_slot().
+
+    Returns:
+      A list of strings.
+    """
+    return self._opt.get_slot_names(*args, **kwargs)
+
+  def get_init_tokens_op(self, num_tokens=-1):
+    """Returns the op to fill the sync_token_queue with the tokens.
+
+    This is supposed to be executed in the beginning of the chief/sync thread
+    so that even if the total_num_replicas is less than replicas_to_aggregate,
+    the model can still proceed as the replicas can compute multiple steps per
+    variable update. Make sure:
+    `num_tokens >= replicas_to_aggregate - total_num_replicas`.
+
+    Args:
+      num_tokens: Number of tokens to add to the queue.
+
+    Returns:
+      An op for the chief/sync replica to fill the token queue.
+
+    Raises:
+      ValueError: If this is called before apply_gradients().
+      ValueError: If num_tokens are smaller than replicas_to_aggregate -
+        total_num_replicas.
+    """
+    if self._gradients_applied is False:
+      raise ValueError(
+          'get_init_tokens_op() should be called after apply_gradients().')
+
+    tokens_needed = self._replicas_to_aggregate - self._total_num_replicas
+    if num_tokens == -1:
+      num_tokens = self._replicas_to_aggregate
+    elif num_tokens < tokens_needed:
+      raise ValueError(
+          'Too few tokens to finish the first step: %d (given) vs %d (needed)' %
+          (num_tokens, tokens_needed))
+
+    if num_tokens > 0:
+      with ops.device(self._global_step.device), ops.name_scope(''):
+        tokens = array_ops.fill([num_tokens], self._global_step)
+        init_tokens = self._sync_token_queue.enqueue_many((tokens,))
+    else:
+      init_tokens = control_flow_ops.no_op(name='no_init_tokens')
+
+    return init_tokens
+
+  def make_session_run_hook(self, is_chief, num_tokens=-1):
+    """Creates a hook to handle SyncReplicasHook ops such as initialization."""
+    return _SyncReplicasOptimizerHook(self, is_chief, num_tokens)
+
+
+class _SyncReplicasOptimizerHook(session_run_hook.SessionRunHook):
+  """A SessionRunHook handles ops related to SyncReplicasOptimizer."""
+
+  def __init__(self, sync_optimizer, is_chief, num_tokens):
+    """Creates hook to handle SyncReplicasOptimizer initialization ops.
+
+    Args:
+      sync_optimizer: `SyncReplicasOptimizer` which this hook will initialize.
+      is_chief: `Bool`, whether is this a chief replica or not.
+      num_tokens: Number of tokens to add to the queue.
+    """
+    self._sync_optimizer = sync_optimizer
+    self._is_chief = is_chief
+    self._num_tokens = num_tokens
+
+  def begin(self):
+    if self._sync_optimizer._gradients_applied is False:  # pylint: disable=protected-access
+      raise ValueError(
+          'SyncReplicasOptimizer.apply_gradient should be called before using '
+          'the hook.')
+    if self._is_chief:
+      self._local_init_op = self._sync_optimizer.chief_init_op
+      self._ready_for_local_init_op = (
+          self._sync_optimizer.ready_for_local_init_op)
+      self._init_tokens_op = self._sync_optimizer.get_init_tokens_op(
+          self._num_tokens)
+    else:
+      self._local_init_op = self._sync_optimizer.local_step_init_op
+      self._ready_for_local_init_op = (
+          self._sync_optimizer.ready_for_local_init_op)
+      self._init_tokens_op = None
+
+  def after_create_session(self, session, coord):
+    """Runs SyncReplicasOptimizer initialization ops."""
+    local_init_success, msg = session_manager._ready(  # pylint: disable=protected-access
+        self._ready_for_local_init_op, session,
+        'Model is not ready for SyncReplicasOptimizer local init.')
+    if not local_init_success:
+      raise RuntimeError(
+          'Init operations did not make model ready for SyncReplicasOptimizer '
+          'local_init. Init op: %s, error: %s' %
+          (self._local_init_op.name, msg))
+    session.run(self._local_init_op)
+    is_closed = session.run(self._sync_optimizer._is_sync_que_closed)
+    assert not is_closed, 'sync_que is closed'
+    if self._init_tokens_op is not None:
+      session.run(self._init_tokens_op)
+
+  def end(self, session):
+    try:
+      is_closed = session.run(self._sync_optimizer._is_sync_que_closed)
+      if not is_closed:
+        logging.info('will close sync token que')
+        session.run(self._sync_optimizer._close_sync_que)
+      else:
+        logging.info('sync token que is closed')
+    except errors_impl.CancelledError:
+      logging.info('sync token que is closed')