(self, session_creator, session_init)
| 483 | |
| 484 | @HIDE_DOC |
| 485 | def initialize(self, session_creator, session_init): |
| 486 | # broadcast_op should be the last setup_graph: it needs to be created |
| 487 | # "right before" the graph is finalized, |
| 488 | # because it needs to capture all the variables (which may be created by callbacks). |
| 489 | self._num_global_variables = len(tf.global_variables()) |
| 490 | self._broadcast_op = self.hvd.broadcast_global_variables(0) |
| 491 | |
| 492 | # it's important that our NewSessionCreator does not finalize the graph |
| 493 | if not isinstance(session_creator, NewSessionCreator): |
| 494 | raise ValueError( |
| 495 | "session_creator has to be `NewSessionCreator` for horovod/byteps training! ") |
| 496 | # NOTE It will fail if GPU was already detected before initializing the session |
| 497 | # https://github.com/tensorflow/tensorflow/issues/8136 |
| 498 | session_creator.config.gpu_options.visible_device_list = str(self._local_rank) |
| 499 | try: |
| 500 | session_creator.config.inter_op_parallelism_threads = mp.cpu_count() // self.hvd.local_size() |
| 501 | except AttributeError: # old horovod does not have local_size |
| 502 | pass |
| 503 | super(HorovodTrainer, self).initialize(session_creator, session_init) |
| 504 | |
| 505 | # This broadcast belongs to the "intialize" stage |
| 506 | # It should not be delayed to the "before_train" stage. |
| 507 | # TODO: |
| 508 | # 1. a allgather helper to concat strings |
| 509 | # 2. check variables on each rank match each other, print warnings, and broadcast the common set. |
| 510 | if self.is_chief: |
| 511 | logger.info("Broadcasting initialization of {} global variables ...".format(self._num_global_variables)) |
| 512 | else: |
| 513 | logger.info("Rank {} waiting for initialization of {} variables ...".format( |
| 514 | self._rank, self._num_global_variables)) |
| 515 | self.sess.run(self._broadcast_op) |
| 516 | |
| 517 | |
| 518 | class BytePSTrainer(HorovodTrainer): |
no test coverage detected