INNER CODE UNIT · Python

register_local_var

horovod/horovod · horovod/_keras/__init__.py:92

        def register_local_var(self, var):
            """Registers a source/variable as worker local. Horovod will not perform any global
            operations on gradients corresponding to these sources and will instead return the local
            gradient."""
            if self._agg_helper:
                self._agg_helper.register_local_var(var)
            elif _IS_TF2:
                self._local_vars.add(var.ref())
            else:
                self._local_vars.add(var)

        def compute_gradients(self, loss, var_list, tape=None):
            return self._compute_gradients(loss, var_list, None, tape)

        def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
            """
            Compute gradients of all trainable variables.

View source record →