(self, n_eval_steps)
Evaluate the model and log metrics.
Evaluate the model and log metrics.
Evaluate the model and log metrics .
def evaluate(self, n_eval_steps):
    """Evaluate the model and log metrics."""
    _, rng = jax_random.split(self._rngs[0])
    # TODO(lukaszkaiser): both model state and parameters by default include
    # the loss layer. Currently, we access the pure-model parameters by just
    # indexing, [0] here. But we should make it more explicit in a better API.
    weights = (self._opt_state.weights[0], self._metrics_weights)
    state = (self._model_state[0], self._metrics_state)
    self.log_step('Evaluation')
    train_eval_slice = itertools.islice(self._train_eval_stream, n_eval_steps)
    train_metrics, _ = self.evaluation_round(train_eval_slice, weights, state,
                                             rng)
    self.log_metrics(train_metrics, self._train_sw, 'train')
    eval_slice = itertools.islice(self._eval_stream, n_eval_steps)
    eval_metrics, _ = self.evaluation_round(eval_slice, weights, state, rng)
    self.log_metrics(eval_metrics, self._eval_sw, 'eval')
    self.log_step('Finished evaluation')

    # Save the learning rate in history.
    self._history.append('train', 'training/learning_rate',
                         self._step, self.learning_rate)
def evaluate ( self , n_eval_steps ) : _ , rng = jax_random . split ( self . _rngs [ 0 ] ) # TODO(lukaszkaiser): both model state and parameters by default include # the loss layer. Currently, we access the pure-model parameters by just # indexing, [0] here. But we should make it more explicit in a better API. weights = ( self . _opt_state . weights [ 0 ] , self . _metrics_weights ) state = ( self . _model_state [ 0 ] , self . _metrics_state ) self . log_step ( 'Evaluation' ) train_eval_slice = itertools . islice ( self . _train_eval_stream , n_eval_steps ) train_metrics , _ = self . evaluation_round ( train_eval_slice , weights , state , rng ) self . log_metrics ( train_metrics , self . _train_sw , 'train' ) eval_slice = itertools . islice ( self . _eval_stream , n_eval_steps ) eval_metrics , _ = self . evaluation_round ( eval_slice , weights , state , rng ) self . log_metrics ( eval_metrics , self . _eval_sw , 'eval' ) self . log_step ( 'Finished evaluation' ) # Save the learning rate in history. self . _history . append ( 'train' , 'training/learning_rate' , self . _step , self . learning_rate )