tensorflow - 使用分布策略在 Estimator 中累积梯度
问题描述
为了减少分布式训练中的同步次数,我想先做梯度的局部累积。这就像您可以拥有多个 GPU,但串行而不是并行。
我想在 estimator.train 循环中使用分布式策略,例如镜像和集体 allreduce 等。
这是我的实现,请给我一些输入:)
首先因为我需要在 session.run() 中运行不同的图表,所以我修改了 estimator.EstimatorSpec 以采取更多的操作。其次,在分布式策略环境中,似乎没有明确的方法在本地 GPU 中创建本地非共享变量。我不得不破解一些 variable_create_scope。
这是被黑的 variable_creator 函数,
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])
我在谷歌官方模型库中的变压器模型上进行了尝试。结果很好。
我的问题是,有没有更好的方法来做到这一点?
我是否应该考虑使用 tf.cond() 来选择 model_fn 中返回的操作,以便不需要修改 Estimator 和 EstimatorSpec?但这似乎非常困难:(
非常感谢!
董
解决方案
我认为您可以通过将 train_ops 传递给估算器来实现这一点。在估计器 model_fn 中单独调用 tensorflow ops 绝对没有效果。因为按照设计,model_fn 在每个训练会话中只调用一次,因此您放入其中的每个操作也将只执行一次。除此之外,所有 tf.cond 分支都将在 model_fn 调用期间被评估和执行。(您可以通过简单的条件日志记录操作来验证此行为。)实现梯度累积的关键是:
- 用 tf.cond 包装所有操作,并结合 tf.no_op 作为 false_fn。
- 让 train_op = tf.group(*accum_ops, [conditional_minimize_op, reset_ops]),但是通过 control_dependencies 控制你的执行顺序,因为 tf.group 并不关心。
- 将你满载的 train_op 传递给 EstimatorSpec
传递给 estimator_spec 或 training_hooks 的那些操作可以在训练过程中动态执行。
这是我的代码,用有限的 GPU 内存微调 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 是一个由 Hook 管理的变量,它会在每个 gradient_accmulation_multiplier 步骤中取值 1,所以这个语句与 tf.math.equal(global_step % gradient_accmulation_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})
推荐阅读
- kdb - KDB:如何获取表中行的不同值?
- javascript - Javascript如何更改svg线性渐变停止颜色
- tensorflow - 调整图像大小时出错:TypeError:无法为
我正在尝试运行一个 Age CGAN 程序,作为其中的一部分,我需要将图像大小从 64x64x3 调整为 192x192x3。在调整大小时,由于某种原因,我收到错误消息
TypeError: Could not build a TypeSpec for <KerasTens
- bash - 在 bash 中,我想将文件名保存在列表中。有什么解决办法吗?
- c# - 使用实体框架代码迁移运行 EnsureScheme 会导致语法错误
- python - 如何在 Python 会话中打开带有重定向的登录页面?
- html - 将html特定表中句子的第一个单词大写
- spring - STS 中的调试错误
- javascript - minikube 集群上的 Vanilla JS ERR_NAME_NOT_RESOLVED
- node.js - 在 forEach nodeJS 之后发送响应