Args: average (bool): whether to average or sum the gradients across processes. compression: `hvd.Compression.fp16` or `hvd.Compression.none`
(self, average=True, compression=None)
| 410 | """ |
| 411 | |
| 412 | def __init__(self, average=True, compression=None): |
| 413 | """ |
| 414 | Args: |
| 415 | average (bool): whether to average or sum the gradients across processes. |
| 416 | compression: `hvd.Compression.fp16` or `hvd.Compression.none` |
| 417 | """ |
| 418 | if 'pyarrow' in sys.modules: |
| 419 | logger.warn("Horovod and pyarrow may conflict due to pyarrow bugs.") |
| 420 | # lazy import |
| 421 | import horovod.tensorflow as hvd |
| 422 | import horovod |
| 423 | hvd_version = tuple(map(int, horovod.__version__.split('.')[:3])) |
| 424 | self.hvd = hvd |
| 425 | |
| 426 | hvd.init() |
| 427 | self.is_chief = hvd.rank() == 0 |
| 428 | self._local_rank = hvd.local_rank() |
| 429 | self._rank = hvd.rank() |
| 430 | self._average = average |
| 431 | self._compression = compression |
| 432 | self._has_compression = hvd_version >= (0, 15, 0) |
| 433 | logger.info("[HorovodTrainer] local rank={}".format(self._local_rank)) |
| 434 | super(HorovodTrainer, self).__init__() |
| 435 | |
| 436 | self.BROADCAST_EVERY_EPOCH = True |
| 437 | |
| 438 | def mpi_enabled(self): |
| 439 | """ |