Args: average (bool): whether to average or sum the gradients across processes.
(self, average=True)
| 533 | does not need to worry about which library is being used under the hood. |
| 534 | """ |
| 535 | def __init__(self, average=True): |
| 536 | """ |
| 537 | Args: |
| 538 | average (bool): whether to average or sum the gradients across processes. |
| 539 | """ |
| 540 | import byteps.tensorflow as bps |
| 541 | self.hvd = bps # BytePS has the same interface as Horovod |
| 542 | self.hvd.allreduce = bps.push_pull # https://github.com/bytedance/byteps/issues/8 |
| 543 | assert os.environ.get("DMLC_ROLE", None) == "worker" |
| 544 | assert "DMLC_WORKER_ID" in os.environ and "DMLC_NUM_WORKER" in os.environ |
| 545 | bps.init() |
| 546 | self.is_chief = bps.rank() == 0 |
| 547 | |
| 548 | self._local_rank = bps.local_rank() |
| 549 | self._rank = bps.rank() |
| 550 | self._average = average |
| 551 | |
| 552 | self._compression = None |
| 553 | self._has_compression = False |
| 554 | logger.info("[BytePSTrainer] local rank={}".format(self._local_rank)) |
| 555 | SingleCostTrainer.__init__(self) |
| 556 | |
| 557 | def mpi_enabled(self): |
| 558 | """ |