Question

假设我有一个TensorFlow变量来跟踪值的平均值。可以使用以下图表摘要更新mean：

mean.assign((step * mean.read() + value) / (step + 1))

不幸的是，这些操作不是原子操作，因此如果图形的两个不同部分尝试更新相同的mean变量，则其中一个更新可能会丢失。

如果我正在跟踪sum，我可以做

sum.assign_add(value, use_locking=True)

一切都会很棒。不幸的是，在其他情况下，可能需要对mean（或std等）进行更复杂的更新，并且可能无法使用tf.assign_add。

问题：有没有办法让第一个代码片段成为原子？

不幸的是我认为答案是否定的，因为（1）我不记得任何这样的机制和（2）我们使优化器C ++操作的原因之一是获得原子行为。我的主要希望来源是XLA，但我不知道是否可以保证这种原子性。

Answer 1

该示例的基本问题是有两个操作 - 读取和后续分配 - 需要以原子方式一起执行。

自2018年初以来，tensorflow团队将 CriticalSection 类添加到codebase。但是，这仅适用于资源变量（正如Geoffrey的评论中所指出的那样）。因此，以下示例中的value需要获取为：

value = tf.get_variable(..., use_resource=True, ...)

虽然我没有对此进行测试，但根据类的文档，原子更新问题应该可解决如下：

def update_mean(step, value):
    old_value = mean.read_value()
    with tf.control_dependencies([old_value]):
        return mean.assign((step * old_value + value) / (step + 1))

cs = tf.CriticalSection()
mean_update = cs.execute(update_mean, step, value)
session.run(mean_update)

本质上，它提供从 execute（）开始到结束的锁定，即覆盖整个赋值操作，包括读取和赋值。

TensorFlow中变量的原子更新

1 个答案: