Накопление градиентов в Оценщике с помощью стратегии распределения

Question

Чтобы уменьшить количество синхронизаций в распределенном обучении, я хочу сначала выполнить локальное накопление градиентов.точно так же, как вы можете иметь несколько графических процессоров, но последовательно, а не параллельно.

Я хочу использовать его в цикле estimator.train со стратегией распределения, такой как зеркальное отображение и коллективное allreduce и т. д.

Вот моя реализация, пожалуйста, дайте мне некоторые входные данные:)

Во-первых, потому что мне нужно запустить другой граф в session.run (), поэтому я изменил estimator.EstimatorSpec, чтобы получить больше операций.Во-вторых, кажется, что нет ясного способа создать локальную переменную без общего доступа в локальном GPU в среде стратегии распределения.Мне пришлось взломать некоторую переменную_создания_области.

Вот взломанная функция переменной_создателя,

def skip_all_scope_variable_creator(next_creator=None, on_device=None, **kwargs):
  #print("skip_all_scope_variable_creator:[{}]".format(kwargs))
  initial_value = kwargs.get("initial_value", None)
  trainable = kwargs.get("trainable", None)
  collections = kwargs.get("collections", None)
  validate_shape = kwargs.get("validate_shape", True)
  caching_device = kwargs.get("caching_device", None)
  name = kwargs.get("name", None)
  variable_def = kwargs.get("variable_def", None)
  dtype = kwargs.get("dtype", None)
  expected_shape = kwargs.get("expected_shape", None)
  import_scope = kwargs.get("import_scope", None)
  constraint = kwargs.get("constraint", None)
  use_resource = kwargs.get("use_resource", None)

  with tf.device(on_device) :
    return resource_variable_ops.ResourceVariable(
      initial_value=initial_value, trainable=trainable,
      collections=collections, validate_shape=validate_shape,
      caching_device=caching_device, name=name, dtype=dtype,
      constraint=constraint, variable_def=variable_def,
      import_scope=import_scope)

Вот мой код внутри model_fn () для создания трех операций,

    loss = loss_from_model
    optimizer = some_optimizer
    tvars = tf.trainable_variables()

    gradients = optimizer.compute_gradients(
      loss, tvars, colocate_gradients_with_ops=True)

    accumulate_pass_num = FLAGS.pass_per_batch

    if accumulate_pass_num > 1 :
      accum_grads = []
      accum_vars = []

      reset_grad_ops = []
      accum_grad_ops = []
      for g,v in gradients:
        accum_vars.append(v)
        if g is not None:
          with tf.variable_creator_scope(lambda next_creator=None, **kwargs: skip_all_scope_variable_creator(next_creator, g.device, **kwargs)):
            print("create accum_grad for variable:{}".format(v.name))
            tmp_grad_on_device = tf.Variable(tf.zeros_like(g), trainable=False, synchronization=tf.VariableSynchronization.ON_READ, collections=[tf.GraphKeys.LOCAL_VARIABLES], name='tmp_accum_grad')
            reset_one_grad_op = tf.assign(tmp_grad_on_device, g, name="reset_accumulated_gradient_op")
            reset_grad_ops.append(reset_one_grad_op)
            # the return of assign_add is the value will be update
            accum_grad_on_device = tmp_grad_on_device.assign_add(g, name="accumulate_gradient")
            accum_grad_ops.append(accum_grad_on_device)
            accum_grads.append(accum_grad_on_device)
        else:
          accum_grads.append(None)

      accumulate_gradients_op = tf.group(*accum_grad_ops, name="grouped_accu_grad_op")
      reset_gradients_op = tf.group(*reset_grad_ops, name="grouped_reset_gradients_op")
      accum_grad_means = [tf.multiply(v, 1.0/accumulate_pass_num) if v is not None else None for v in accum_grads]
      accum_grads_vars = zip(accum_grad_means, accum_vars)
      minimize_op = optimizer.apply_gradients(
        accum_grads_vars, global_step=global_step, name="train")

     update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
     train_op = tf.group(minimize_op, update_ops)
     return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op, accumulate_gradients_op=accumulate_gradients_op, reset_gradients_op=reset_gradients_op, accumulate_pass_num=accumulate_pass_num)

Вот модифицированный estimator.train () для запуска разных операций,

      while not mon_sess.should_stop():
        if estimator_spec.accumulate_pass_num > 1 :
          # reset gradiends first
          mon_sess.run([estimator_spec.reset_gradients_op])
          for _ in range(estimator_spec.accumulate_pass_num-2):
            mon_sess.run([estimator_spec.accumulate_gradients_op])

        _, loss = mon_sess.run([estimator_spec.train_op, estimator_spec.loss])

Я попробовал это на модели трансформатора в официальном репозитории моделей Google.Результаты были хорошими.

Мой вопрос: есть ли лучшие способы сделать это?

Должен ли я рассмотреть использование tf.cond () для выбора операций, возвращаемых в model_fn, чтобы Estimator и EstimatorSpec не нуждались в изменении?Но это кажется очень сложным: (

Большое спасибо!

Донг

Answer 1

Я думаю, что вы можете достичь этого, передавая train_ops в оценщик.Вызов одних только тензорных операций внутри оценщика model_fn не дает абсолютно никакого эффекта.Потому что по замыслу model_fn вызывается только один раз за тренировку, следовательно, каждая операция, которую вы в нее вставляете, также будет выполняться только один раз.Кроме того, все ветви tf.cond будут оцениваться и выполняться во время вызова model_fn. (Вы можете проверить это поведение с помощью простой операции условного ведения журнала.) Ключ для достижения накопления градиента:

1. Оберните все ваши операции с помощью tf.cond, в сочетании с tf.no_op как false_fn.
2. Пусть train_op = tf.group (* gast_ops, [conditional_minimize_op, reset_ops]), но контролируют ваш порядок выполнения с помощью control_dependencies, потому чтоtf.group не заботится.
3. Передайте полностью загруженный train_op EstimatorSpec

Эти операции, переданные estimator_spec или training_hooks, могут выполняться динамически во время процесса обучения.

Вот мой код, настраивающий BERT с ограниченной памятью графического процессора:

# compute batch gradient
grads = tf.gradients(loss, tvars)
(grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)
# this is a list of sum(dy/dx) for each variable that must be paired with a tvars list.
# element may be an IndexedSlices object that does not support assignning, e.g. [g.assign(value) for g in grads]
# some of the elements are None, meaning y and x does not depend on each other.
# Nonetypes must be handled using Python, tensorflow cannot convert Nonetypes to 0.

# declare a temp variable for summation
sum_gradient = [tf.get_variable(name="sum_grads" + str(i), shape=tv.shape,
                                initializer=tf.zeros_initializer,
                                trainable=False,
                                dtype=tf.float32,
                                collections=[tf.GraphKeys.LOCAL_VARIABLES]) for i, tv in enumerate(tvars)]
sum_ops = []
unused_variable_in_batch = []

# gradient accumulation
for i, gv in enumerate(grads):
    if gv is not None:
        sum_ops.append(sum_gradient[i].assign_add(gv, name="accumulate_gradient"))
    else:
        unused_variable_in_batch.append(sum_gradient[i])
        sum_gradient[i] = None

# NOTE : calling .assign_add does NOTHING in estimator, must wrap them all and handle them via train_ops

def apply_accumulated_gradients(sums):
    # normalize gradient
    normalize_ops = []
    for i, g in enumerate(sums):
        if g is not None:
            normalize_ops.append(sums[i].assign(tf.multiply(g, 1 / gradient_accmulation_multiplier)))
            # assign to make sure it still is a variable, or else it will become a Tensor
    with tf.control_dependencies(normalize_ops):
        minimize_op = optimizer.apply_gradients(zip(sums, tvars), global_step=global_step)
    return tf.group(minimize_op, *normalize_ops, name="apply_accumulated_gradients")

train_op = tf.cond(tf.math.equal(global_step % gradient_accmulation_multiplier, 0),
                   lambda: apply_accumulated_gradients(sum_gradient),
                   lambda: optimizer.apply_gradients(zip([None for _ in grads], tvars), global_step=global_step))

# reset accumulation when necessary
def reset():
    counter = 0
    for i, s in enumerate(sum_gradient):
        if s is None:
            # restore reference from None to the original variable
            sum_gradient[i] = unused_variable_in_batch[counter]
            counter += 1
    return tf.group([s.assign(tf.zeros_like(s)) for s in sum_gradient])

with tf.control_dependencies([train_op]):
    reset_ops = tf.cond(tf.math.equal(do_update, 1.),
                        reset,
                        tf.no_op)
# the 2 branches must have identical structure, [op1, op2, ...] || no_op cannot be valid cond branch.
# tf.group to convert all resets into 1 op and match with no_op: tf.group() || np_op

# Increment global step
new_global_step = global_step + 1
train_op = tf.group(*sum_ops, [train_op, global_step.assign(new_global_step), reset_ops])

logging_hook = tf.train.LoggingTensorHook({"accuracy": "acc"},                                                          
                                    every_n_iter=gradient_accmulation_multiplier)
output_spec = tf.estimator.EstimatorSpec(
                mode=mode,
                loss=loss,
                train_op=train_op,
                training_hooks=[logging_hook, accumulation_hook]  # wrap with a list
            )

Я применил отсечение к пакетным градиентам и просто взял среднее из них.Этот подход сработал для меня, но я рекомендую вам внимательно следить за поведением потерь в вашем наборе данных.

Кроме того, о tf.cond (tf.math.equal (do_update, 1.), ..., ...), do_update - это переменная, управляемая хуком, она будет принимать значение 1 для каждого шага градиента_аксмуляции_множителя, поэтому этот оператор имеет тот же эффект, что и tf.math.equal (global_step% градиент_аксмуляция_multiplier, 0).Это просто другой способ.

Код для Hook выглядит следующим образом:

class GradientAccumulationHook(session_run_hook.SessionRunHook):
    """
    Puts a certain tf.Variable to 1 once every certain steps.
    """

    def __init__(self, frequency, variable):
        self._step = 0
        self._flag = 0.
        self._freq = frequency
        self._input_placeholder = tf.placeholder(tf.float32)
        self.assign_op = variable.assign(self._input_placeholder)

    def begin(self):
        # a hook can modify graph at begin(), after this the graph will be finalized
        self._step = tf.train.get_global_step()

    def before_run(self, run_context):
        step = run_context.session.run(self._step)  # evaluate tensor to get a step number
        self._flag = 1. if step % self._freq == 0 and step != 0 else 0.
        run_context.session.run(self.assign_op, feed_dict={self._input_placeholder: self._flag})