(self, output_dir, init_checkpoint=None)
Reset the model parameters.

    Restores the parameters from the given output_dir if a checkpoint exists,
    otherwise randomly initializes them.

    Does not re-jit the model.

    Args:
      output_dir: Output directory.
      init_checkpoint: Initial checkpoint (default $output_dir/model.pkl.gz)
Reset the model parameters.
Reset the model parameters .
def reset(self, output_dir, init_checkpoint=None):
    """Reset the model parameters.

    Restores the parameters from the given output_dir if a checkpoint exists,
    otherwise randomly initializes them.

    Does not re-jit the model.

    Args:
      output_dir: Output directory.
      init_checkpoint: Initial checkpoint (default $output_dir/model.pkl.gz)
    """
    self.close()
    self._output_dir = output_dir
    if output_dir is not None:
      tf.io.gfile.makedirs(output_dir)
    else:
      assert not self._should_save_checkpoints
      assert not self._should_write_summaries

    # Create summary writers and history.
    if self._should_write_summaries:
      self._train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, 'train'),
                                              enable=self._is_chief)
      self._eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, 'eval'),
                                             enable=self._is_chief)

    # Reset the train and eval streams.
    self._train_stream = _repeat_stream(self._inputs.train_stream,
                                        self._n_devices)
    # TODO(lukaszkaiser): add an option to evaluate exactly on the full eval
    #   set by adding a padding and stopping the stream when too large.
    self._eval_stream = _repeat_stream(
        self._inputs.eval_stream, self._n_devices)
    self._train_eval_stream = _repeat_stream(
        self._inputs.train_eval_stream, self._n_devices)

    # Restore the training state.
    if output_dir is not None:
      state = load_trainer_state(output_dir, self._model_with_loss,
                                 init_checkpoint)
    else:
      state = TrainerState(step=None, opt_state=None,
                           history=trax_history.History(), model_state=None)
    self._step = state.step or 0
    history = state.history
    self._history = history
    if state.opt_state:
      opt_state = state.opt_state
      model_state = state.model_state
    else:
      opt_state, model_state = self._new_opt_state_and_model_state()
      model_state = self._for_n_devices(model_state)
    self._opt_state = OptState(*self._for_n_devices(opt_state))
    self._model_state = model_state
    if not state.opt_state and self._should_save_checkpoints:
      self.save_state(keep=False)
def reset ( self , output_dir , init_checkpoint = None ) : self . close ( ) self . _output_dir = output_dir if output_dir is not None : tf . io . gfile . makedirs ( output_dir ) else : assert not self . _should_save_checkpoints assert not self . _should_write_summaries # Create summary writers and history. if self . _should_write_summaries : self . _train_sw = jaxboard . SummaryWriter ( os . path . join ( output_dir , 'train' ) , enable = self . _is_chief ) self . _eval_sw = jaxboard . SummaryWriter ( os . path . join ( output_dir , 'eval' ) , enable = self . _is_chief ) # Reset the train and eval streams. self . _train_stream = _repeat_stream ( self . _inputs . train_stream , self . _n_devices ) # TODO(lukaszkaiser): add an option to evaluate exactly on the full eval #   set by adding a padding and stopping the stream when too large. self . _eval_stream = _repeat_stream ( self . _inputs . eval_stream , self . _n_devices ) self . _train_eval_stream = _repeat_stream ( self . _inputs . train_eval_stream , self . _n_devices ) # Restore the training state. if output_dir is not None : state = load_trainer_state ( output_dir , self . _model_with_loss , init_checkpoint ) else : state = TrainerState ( step = None , opt_state = None , history = trax_history . History ( ) , model_state = None ) self . _step = state . step or 0 history = state . history self . _history = history if state . opt_state : opt_state = state . opt_state model_state = state . model_state else : opt_state , model_state = self . _new_opt_state_and_model_state ( ) model_state = self . _for_n_devices ( model_state ) self . _opt_state = OptState ( * self . _for_n_devices ( opt_state ) ) self . _model_state = model_state if not state . opt_state and self . _should_save_checkpoints : self . save_state ( keep = False )