def save(self, checkpoint_dir=None) -> str: """Saves the current model state to a checkpoint. Subclasses should override ``save_checkpoint()`` instead to save state. This method dumps additional metadata alongside the saved path. If a remote checkpoint dir is given, this will also sync up to remote storage. Args: checkpoint_dir (str): Optional dir to place the checkpoint. Returns: str: path that points to xxx.pkl file. Note the return path should match up with what is expected of `restore()`. """ checkpoint_dir = TrainableUtil.make_checkpoint_dir( checkpoint_dir or self.logdir, index=self.iteration) checkpoint = self.save_checkpoint(checkpoint_dir) trainable_state = self.get_state() checkpoint_path = TrainableUtil.process_checkpoint( checkpoint, parent_dir=checkpoint_dir, trainable_state=trainable_state) self._postprocess_checkpoint(checkpoint_dir) # Maybe sync to cloud self._maybe_save_to_cloud(checkpoint_dir) return checkpoint_path
def save(self, checkpoint_dir=None): """Saves the current model state to a checkpoint. Subclasses should override ``save_checkpoint()`` instead to save state. This method dumps additional metadata alongside the saved path. If a remote checkpoint dir is given, this will also sync up to remote storage. Args: checkpoint_dir (str): Optional dir to place the checkpoint. Returns: str: Checkpoint path or prefix that may be passed to restore(). """ checkpoint_dir = TrainableUtil.make_checkpoint_dir( checkpoint_dir or self.logdir, index=self.iteration) checkpoint = self.save_checkpoint(checkpoint_dir) trainable_state = self.get_state() checkpoint_path = TrainableUtil.process_checkpoint( checkpoint, parent_dir=checkpoint_dir, trainable_state=trainable_state) # Maybe sync to cloud self._maybe_save_to_cloud() return checkpoint_path
def write_checkpoint(trial: Trial, index: int): checkpoint_dir = TrainableUtil.make_checkpoint_dir(trial.logdir, index=index) result = {"training_iteration": index} with open(os.path.join(checkpoint_dir, "cp.json"), "w") as f: json.dump(result, f) tune_cp = _TuneCheckpoint(_TuneCheckpoint.PERSISTENT, checkpoint_dir, result) trial.saving_to = tune_cp trial.on_checkpoint(tune_cp) return checkpoint_dir
def write_checkpoint(trial: Trial, index: int): checkpoint_dir = TrainableUtil.make_checkpoint_dir( trial.logdir, index=index ) result = {"training_iteration": index} with open(os.path.join(checkpoint_dir, "cp.json"), "w") as f: json.dump(result, f) tune_cp = _TrackedCheckpoint( dir_or_data=checkpoint_dir, storage_mode=CheckpointStorage.PERSISTENT, metrics=result, ) trial.saving_to = tune_cp return checkpoint_dir
def setUp(self): self.checkpoint_dir = os.path.join( ray._private.utils.get_user_temp_dir(), "tune", "MyTrainable123") self.checkpoint_dir = TrainableUtil.make_checkpoint_dir( self.checkpoint_dir, "0")