def build(self, get_grad_fn, get_opt_fn):
        """
        Args:
            get_grad_fn (-> [(grad, var)]):
            get_opt_fn (-> tf.train.Optimizer): callable which returns an optimizer

        Returns:
            (tf.Operation, tf.Operation, tf.Operation):

            1. the training op.

            2. the op which sync all the local variables from PS.
            This op should be run before training.

            3. the op which sync all the local `MODEL_VARIABLES` from PS.
            You can choose how often to run it by yourself.
        """
        with override_to_local_variable():
            get_global_step_var()

        get_opt_fn = memoized(get_opt_fn)
        # Build the optimizer first, before entering any tower.
        # This makes sure that learning_rate is a global variable (what we expect)
        get_opt_fn()    # TODO get_opt_fn called before main graph was built

        # Ngpu * Nvar * 2
        grad_list = DataParallelBuilder.build_on_towers(
            self.towers, get_grad_fn,
            devices=self.raw_devices,
            use_vs=[True] * len(self.towers))  # open vs at each tower
        DataParallelBuilder._check_grad_list(grad_list)

        avg_grads = aggregate_grads(
            grad_list, colocation=False, devices=self.raw_devices)
        with tf.device(self.param_server_device):
            ps_var_grads = DistributedReplicatedBuilder._apply_shadow_vars(avg_grads)
            var_update_ops = self._apply_gradients_and_copy(
                get_opt_fn(), grad_list, ps_var_grads)
            self._shadow_vars = [v for (__, v) in ps_var_grads]
            self._shadow_model_vars = DistributedReplicatedBuilder._shadow_model_variables(self._shadow_vars)

        # TODO add options to synchronize less
        main_fetch = tf.group(*var_update_ops, name='main_fetches')
        train_op = self._add_sync_queues_and_barrier(
            'post_copy_barrier', [main_fetch])

        # initial local_vars syncing
        with tf.name_scope('initial_sync_variables'):
            initial_sync_op = self._get_initial_sync_op()
        if len(self._shadow_model_vars) and self.is_chief:
            with tf.name_scope('sync_model_variables'):
                model_sync_op = self._get_sync_model_vars_op()
        else:
            model_sync_op = None
        return train_op, initial_sync_op, model_sync_op
def build ( self , get_grad_fn , get_opt_fn ) : with override_to_local_variable ( ) : get_global_step_var ( ) get_opt_fn = memoized ( get_opt_fn ) get_opt_fn ( ) grad_list = DataParallelBuilder . build_on_towers ( self . towers , get_grad_fn , devices = self . raw_devices , use_vs = [ True ] * len ( self . towers ) ) DataParallelBuilder . _check_grad_list ( grad_list ) avg_grads = aggregate_grads ( grad_list , colocation = False , devices = self . raw_devices ) with tf . device ( self . param_server_device ) : ps_var_grads = DistributedReplicatedBuilder . _apply_shadow_vars ( avg_grads ) var_update_ops = self . _apply_gradients_and_copy ( get_opt_fn ( ) , grad_list , ps_var_grads ) self . _shadow_vars = [ v for ( __ , v ) in ps_var_grads ] self . _shadow_model_vars = DistributedReplicatedBuilder . _shadow_model_variables ( self . _shadow_vars ) main_fetch = tf . group ( * var_update_ops , name = 'main_fetches' ) train_op = self . _add_sync_queues_and_barrier ( 'post_copy_barrier' , [ main_fetch ] ) with tf . name_scope ( 'initial_sync_variables' ) : initial_sync_op = self . _get_initial_sync_op ( ) if len ( self . _shadow_model_vars ) and self . is_chief : with tf . name_scope ( 'sync_model_variables' ) : model_sync_op = self . _get_sync_model_vars_op ( ) else : model_sync_op = None return train_op , initial_sync_op , model_sync_op
Args:
            get_grad_fn (-> [(grad, var)]):
            get_opt_fn (-> tf.train.Optimizer): callable which returns an optimizer

        Returns:
            (tf.Operation, tf.Operation, tf.Operation):

            1. the training op.

            2. the op which sync all the local variables from PS.
            This op should be run before training.

            3. the op which sync all the local `MODEL_VARIABLES` from PS.
            You can choose how often to run it by yourself.