Args: grads (list): list of gradient tensors Returns: packed list of gradient tensors to be aggregated.
(self, grads)
| 419 | return True |
| 420 | |
| 421 | def pack(self, grads): |
| 422 | """ |
| 423 | Args: |
| 424 | grads (list): list of gradient tensors |
| 425 | |
| 426 | Returns: |
| 427 | packed list of gradient tensors to be aggregated. |
| 428 | """ |
| 429 | for i, g in enumerate(grads): |
| 430 | assert g.shape == self._shapes[i] |
| 431 | |
| 432 | with cached_name_scope("GradientPacker", top_level=False): |
| 433 | concat_grads = tf.concat([tf.reshape(g, [-1]) for g in grads], 0, name='concatenated_grads') |
| 434 | # concat_grads = tf.cast(concat_grads, tf.float16) |
| 435 | grad_packs = tf.split(concat_grads, self._split_sizes) |
| 436 | return grad_packs |
| 437 | |
| 438 | def unpack(self, grad_packs): |
| 439 | with cached_name_scope("GradientPacker", top_level=False): |
no test coverage detected