Represents a reduction method.
Reduces the variable updates by averaging them.
Reduces the variable updates by summing them.
Mirrors value
to all worker devices.
Mirrors value
to all worker devices.
Value to broadcast.
Destination devices.
Mirrored value.
Executes block
within a scope that controls which devices variables will be created on.
Executes block
within a scope that controls which devices variables will be created on.
No operations should be added to the graph inside this scope; it should only be used when creating variables (some
implementations work by changing variable creation and others work by using a colocateWith
scope). This may only
be used inside DistributionStrategy.scope
.
For example:
distributionStrategy.scope { val variable1 = tf.variable(...) distributionStrategy.colocateVariablesWith(Set(variable1.op)) { // `variable2` and `variable3` will be created on the same device(s) as `variable1`. val variable2 = tf.variable(...) val variable3 = tf.variable(...) } def fn(v1: Variable, v2: Variable, v3: Variable): Unit = { // Operates on `v1` from `variable1`, `v2` from `variable2`, and `v3` from `variable3`. } // `fn` runs on every device `v1` is on, and `v2` and `v3` will be there too. distributionStrategy.update(variable1, fn, variable2, variable3) }
Variables created in block
will be on the same set of devices as these ops.
Code block to execute in this scope.
Value returned by block
.
Returns the current device.
Returns the current device.
Returns the current distribution strategy.
Returns the current distribution strategy.
Returns the current device if in a distributionStrategy.update()
call.
Returns the current device if in a distributionStrategy.update()
call.
Runs fn
once per tower.
Runs fn
once per tower.
fn
may call tf.currentTowerContext
to access fields and methods such as towerID
and mergeCall()
.
mergeCall()
is used to communicate between the towers and re-enter the cross-tower context. All towers pause
their execution having encountered a mergeCall()
call. After that the mergeFn
-function is executed. Its
results are then unwrapped and given back to each tower call. After that execution resumes until fn
is complete
or another mergeCall()
is encountered.
For example:
// Called once in "cross-tower" context. def mergeFn(distributionStrategy: DistributionStrategy, threePlusTowerID: Int): tf.Output = { // Sum the values across towers. tf.addN(distribution.unwrap(threePlusTowerID)) } // Called once per tower in `distributionStrategy`, in a "tower" context. def fn(three: Int): Output = { val towerContext = tf.currentTowerContext val v = three + towerContext.towerID // Computes the sum of the `v` values across all towers. val s = towerContext.mergeCall(mergeFn(_, v)) s + v } distributionStrategy.scope { // In "cross-tower" context ... val mergedResults = distributionStrategy.forEachTower(() => fn(3)) // `mergedResults` has the values from every tower execution of `fn`. val resultsList = distributionStrategy.unwrap(mergedResults) }
Function that will be run once per tower.
Wrapped values that will be unwrapped when invoking fn
on each tower.
Merged return value of fn
across all towers.
Executes block
within a scope where new variables will not be mirrored.
Executes block
within a scope where new variables will not be mirrored.
There will still be one component variable per tower, but there is no requirement that they stay in sync. Instead,
when saving them or calling fetch()
, we use the value that results when calling reduce()
on all the towers'
variables. Note that tower-local implies not trainable. Instead, it is expected that each tower will directly
update (e.g., using assignAdd()
) its local variable instance but only the aggregated value (accessible using
fetch()
) will be exported from the model. When it is acceptable to only aggregate on export, we greatly reduce
communication overhead by using tower-local variables.
Note that all component variables will be initialized to the same value, using the initialization expression from the first tower. The values will match even if the initialization expression uses random numbers.
Reduction method used to get the value to save when creating checkpoints.
Code block to execute in this scope.
Value returned by block
.