Python tf.distribute.StrategyExtended.update用法及代码示例

用法

update(
    var, fn, args=(), kwargs=None, group=True
)

运行fn 以使用镜像到相同设备的输入更新var。

tf.distribute.StrategyExtended.update 采用要更新的分布式变量 var、更新函数 fn 以及 args 和 kwargs 用于 fn 。它将 fn 应用于 var 的每个组件变量，并从 args 和 kwargs 传递相应的值。 args 和 kwargs 都不能包含 per-replica 值。如果它们包含镜像值，它们将在调用 fn 之前被解包。例如，fn 可以是 assign_add 并且 args 可以是镜像 DistributedValues，其中每个组件都包含要添加到此镜像变量 var 的值。调用update 将在var 的每个组件变量上调用assign_add，并在该设备上使用相应的张量值。

示例用法：

strategy = tf.distribute.MirroredStrategy(['GPU:0', 'GPU:1']) # With 2
devices
with strategy.scope():
  v = tf.Variable(5.0, aggregation=tf.VariableAggregation.SUM)
def update_fn(v):
  return v.assign(1.0)
result = strategy.extended.update(v, update_fn)
# result is
# Mirrored:{
#  0:tf.Tensor(1.0, shape=(), dtype=float32),
#  1:tf.Tensor(1.0, shape=(), dtype=float32)
# }

如果var跨多个设备镜像，则该方法实现逻辑如下：

results = {}
for device, v in var:
  with tf.device(device):
    # args and kwargs will be unwrapped if they are mirrored.
    results[device] = fn(v, *args, **kwargs)
return merged(results)

否则，此方法返回与 var 并置的 fn(var, *args, **kwargs) 。