A version of `apply_gradients` for cross-replica context. This is a version of `apply_gradients()` for when you are using a `DistributionStrategy` and are in a cross-replica context. If in a replica context, use `apply_gradients()` as normal. Args: distribution: A `Distributi
(self,
distribution,
grads_and_vars,
global_step=None,
name=None)
| 746 | return apply_fn() |
| 747 | |
| 748 | def _distributed_apply(self, |
| 749 | distribution, |
| 750 | grads_and_vars, |
| 751 | global_step=None, |
| 752 | name=None): |
| 753 | """A version of `apply_gradients` for cross-replica context. |
| 754 | |
| 755 | This is a version of `apply_gradients()` for when you are using a |
| 756 | `DistributionStrategy` and are in a cross-replica context. If in a |
| 757 | replica context, use `apply_gradients()` as normal. |
| 758 | |
| 759 | Args: |
| 760 | distribution: A `DistributionStrategy` object. |
| 761 | grads_and_vars: List of (gradient, variable) pairs as returned by |
| 762 | `compute_gradients()`, and then aggregated across replicas. |
| 763 | global_step: Optional (mirrored) `Variable` to increment by one |
| 764 | after the variables have been updated. |
| 765 | name: Optional name for the returned operation. Default to the |
| 766 | name passed to the `Optimizer` constructor. |
| 767 | |
| 768 | Returns: |
| 769 | An `Operation` that applies the specified gradients across all |
| 770 | replicas. If `global_step` was not None, that operation also |
| 771 | increments `global_step` |
| 772 | """ |
| 773 | name = name if name is not None else self.get_name() |
| 774 | def apply_fn(): |
| 775 | reduced_grads = distribution.extended.batch_reduce_to( |
| 776 | ds_reduce_util.ReduceOp.SUM, grads_and_vars) |
| 777 | var_list = [v for _, v in grads_and_vars] |
| 778 | rgrads_and_vars = zip(reduced_grads, var_list) |
| 779 | |
| 780 | # Note that this is called in a cross-replica context. |
| 781 | with ops.init_scope(): |
| 782 | self._create_slots(var_list) |
| 783 | |
| 784 | def update(v, g): |
| 785 | """Apply gradients to a replica variable.""" |
| 786 | assert v is not None |
| 787 | |
| 788 | try: |
| 789 | # Convert the grad to Tensor or IndexedSlices if necessary. |
| 790 | g = ops.convert_to_tensor_or_indexed_slices(g) |
| 791 | except TypeError: |
| 792 | raise TypeError("Gradient must be convertible to a Tensor" |
| 793 | " or IndexedSlices, or None: %s" % g) |
| 794 | if not isinstance(g, (ops.Tensor, ops.IndexedSlices)): |
| 795 | raise TypeError( |
| 796 | "Gradient must be a Tensor, IndexedSlices, or None: %s" % g) |
| 797 | p = _get_processor(v) |
| 798 | |
| 799 | if context.executing_eagerly() or ( |
| 800 | resource_variable_ops.is_resource_variable(v) and |
| 801 | not v._in_graph_mode): # pylint: disable=protected-access |
| 802 | scope_name = v.name.split(":")[0] |
| 803 | else: |
| 804 | scope_name = v.op.name |
| 805 |