Python tf.compat.v1.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) 。