diff --git a/astrai/serialization/checkpoint.py b/astrai/serialization/checkpoint.py index 37ef04c..ea5a9cb 100644 --- a/astrai/serialization/checkpoint.py +++ b/astrai/serialization/checkpoint.py @@ -2,7 +2,6 @@ import io import json -import os import time from dataclasses import dataclass, field from pathlib import Path @@ -178,6 +177,7 @@ class Checkpoint: epoch=meta.get("epoch", 0), consumed_samples=meta.get("consumed_samples", 0), extra=extra, + meta=meta, config=config, ) diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 8f36896..c33bf1b 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -54,10 +54,12 @@ class TrainContextBuilder: config: TrainConfig, ): self.config = config - self._resume_dir: Optional[str] = None + self._param_path: Optional[str] = None + self._resume: bool = False - def with_resume_dir(self, resume_dir: Optional[str]) -> Self: - self._resume_dir = resume_dir + def with_param_path(self, param_path: Optional[str], resume: bool = False) -> Self: + self._param_path = param_path + self._resume = resume return self def build(self) -> TrainContext: @@ -74,8 +76,8 @@ class TrainContextBuilder: model = model.to(device=device) model_config = {} - if self._resume_dir: - config_path = Path(self._resume_dir) / "config.json" + if self._param_path: + config_path = Path(self._param_path) / "config.json" if config_path.exists(): model_config = load_json(config_path) @@ -91,23 +93,29 @@ class TrainContextBuilder: executor=executor, ) - if self._resume_dir: - checkpoint = Checkpoint.load_any(self._resume_dir) + if self._param_path: + checkpoint = Checkpoint.load_any(self._param_path) if checkpoint is not None: model.load_state_dict(checkpoint.state_dict, strict=False) if checkpoint.config: context.model_config = checkpoint.config - context.epoch = checkpoint.epoch or cfg.start_epoch - if checkpoint.consumed_samples > 0: - per_step = ( - cfg.batch_per_device * context.world_size * cfg.grad_accum_steps - ) - context.consumed_samples = ( - checkpoint.consumed_samples // per_step - ) * per_step - else: - context.consumed_samples = cfg.start_samples * context.world_size - context.checkpoint = checkpoint + + if self._resume: + context.epoch = checkpoint.epoch or cfg.start_epoch + if checkpoint.consumed_samples > 0: + per_step = ( + cfg.batch_per_device + * context.world_size + * cfg.grad_accum_steps + ) + context.consumed_samples = ( + checkpoint.consumed_samples // per_step + ) * per_step + else: + context.consumed_samples = ( + cfg.start_samples * context.world_size + ) + context.checkpoint = checkpoint if cfg.lora is not None: inject_lora( diff --git a/astrai/trainer/trainer.py b/astrai/trainer/trainer.py index a8fc3c9..2e1dcea 100644 --- a/astrai/trainer/trainer.py +++ b/astrai/trainer/trainer.py @@ -52,9 +52,11 @@ class Trainer: if method: method(context) - def _trainer_loop(self, resume_dir: Optional[str] = None): + def _trainer_loop(self, param_path: Optional[str] = None, resume: bool = False): context = ( - TrainContextBuilder(self.train_config).with_resume_dir(resume_dir).build() + TrainContextBuilder(self.train_config) + .with_param_path(param_path, resume=resume) + .build() ) executor = context.executor self._call_callbacks("on_train_begin", context) @@ -95,7 +97,7 @@ class Trainer: finally: self._call_callbacks("on_train_end", context) - def train(self, resume_dir: Optional[str] = None): + def train(self, param_path: Optional[str] = None, resume: bool = False): cfg = self.train_config spawn_parallel_fn( self._trainer_loop, @@ -105,5 +107,6 @@ class Trainer: master_port=cfg.master_port, device_type=cfg.device_type, start_method=cfg.start_method, - resume_dir=resume_dir, + param_path=param_path, + resume=resume, ) diff --git a/scripts/tools/train.py b/scripts/tools/train.py index af3a70f..767f8a3 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -116,6 +116,13 @@ def parse_args() -> argparse.Namespace: required=True, help="Path to the model parameters or resume checkpoint.", ) + parser.add_argument( + "--resume", + action="store_true", + default=False, + help="Resume training from checkpoint at --param_path " + "(restore epoch, consumed_samples, optimizer & scheduler state).", + ) parser.add_argument( "--n_epoch", type=int, default=1, help="Number of epochs to train." @@ -385,6 +392,7 @@ def train( train_type: str, param_path: str, data_root_path: str, + resume: bool, n_epoch: int, batch_per_device: int, start_epoch: int, @@ -531,7 +539,7 @@ def train( ) trainer = Trainer(train_config) - trainer.train(resume_dir=param_path) + trainer.train(param_path=param_path, resume=resume) if __name__ == "__main__": diff --git a/tests/trainer/test_early_stopping.py b/tests/trainer/test_early_stopping.py index 81bd490..2dd72ad 100644 --- a/tests/trainer/test_early_stopping.py +++ b/tests/trainer/test_early_stopping.py @@ -46,7 +46,7 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset): # Resume from latest checkpoint load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_step_1") trainer = Trainer(train_config) - trainer.train(resume_dir=load_dir) + trainer.train(param_path=load_dir, resume=True) # Verify checkpoint was saved at expected step load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_step_5")