def _broadcast_model(model, optimizer, backend): if _executing_eagerly(): # TensorFlow 2.0 or TensorFlow eager broadcast_variables(model.variables, root_rank=0) broadcast_variables(optimizer.variables(), root_rank=0) else: bcast_op = broadcast_variables(_global_variables(), root_rank=0) backend.get_session().run(bcast_op)
def __init__(self, variables=None, session=None, **kwargs): self.variables = variables or _global_variables() self.session = session or _default_session() self._bcast_op = broadcast_variables(self.variables, root_rank=0) self._eval_fn = self._to_numpy if _executing_eagerly() else self._eval_var self._assign_fn = self._assign_var if _IS_TF2 else self._load_var self._save_model() bcast_obj = broadcast_object_fn(session=session) if not _executing_eagerly() else broadcast_object def broadcast_object_with_session(obj): return bcast_obj(obj) super(TensorFlowState, self).__init__(bcast_object=broadcast_object_with_session, get_rank=rank, **kwargs)
def broadcast_global_variables(root_rank): """Broadcasts all global variables from root rank to all other processes. **NOTE:** deprecated in TensorFlow 2.0. Arguments: root_rank: rank of the process from which global variables will be broadcasted to all other processes. """ if _executing_eagerly(): raise RuntimeError( "hvd.broadcast_global_variables() does not support eager execution. " "Please use `hvd.broadcast_variables(<model/optimizer variables>)` instead." ) return broadcast_variables(_global_variables(), root_rank)
def __init__(self, model, optimizer=None, backend=None, **kwargs): self.model = model if not _model_built(model): raise ValueError('Model must be built first. Run `model.build(input_shape)`.') self.optimizer = optimizer or model.optimizer self.backend = backend self._save_model() if not backend or _executing_eagerly(): self._bcast_model = lambda: _broadcast_model(self.model, self.optimizer, backend=self.backend) bcast_object = broadcast_object else: # For TensorFlow v1, we need to reuse the broadcast op to prevent incrementing the uids bcast_op = broadcast_variables(_global_variables(), root_rank=0) self._bcast_model = lambda: self.backend.get_session().run(bcast_op) bcast_object = broadcast_object_fn(session=self.backend.get_session()) super(TensorFlowKerasState, self).__init__(bcast_object=bcast_object, get_rank=rank, **kwargs)
def _broadcast_model(self): broadcast_variables(self._model.variables, root_rank=0) broadcast_variables(self._optimizer.variables(), root_rank=0)
def set_broadcast_variables(self, variables): if self._bcast_op is None: self._variables = variables self._bcast_op = broadcast_variables(self._variables, root_rank=0)