INNER CODE UNIT · Python
_aggregate_gradients
bytedance/byteps · byteps/_keras/__init__.py:43
def _aggregate_gradients(self, grads_and_vars):
gradients = [grad for grad, var in grads_and_vars]
return self._push_pull(gradients)
def _push_pull(self, gradients):
self._aggregated_gradients = True
if bps.size() > 1:
averaged_gradients = []
with tf.name_scope(self._name + "_Push_Pull") as scope:
for grad in gradients:
if grad is not None:
if self._sparse_as_dense and \
isinstance(grad, tf.IndexedSlices):
grad = tf.convert_to_tensor(grad)
avg_grad = bps.push_pull(grad, scope,
device_dense=self._device_dense,
device_sparse=self._device_sparse,
compression=self._compression)