diff --git a/docs/getting_started.md b/docs/getting_started.md index e0a1cae..3098ea1 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -241,11 +241,12 @@ Optional arguments are: - `--no-validate` (**not suggested**): By default, the codebase will perform evaluation at every k iterations during the training. To disable this behavior, use `--no-validate`. - `--work-dir ${WORK_DIR}`: Override the working directory specified in the config file. -- `--resume-from ${CHECKPOINT_FILE}`: Resume from a previous checkpoint file. +- `--resume-from ${CHECKPOINT_FILE}`: Resume from a previous checkpoint file (to continue the training process). +- `--load-from ${CHECKPOINT_FILE}`: Load weights from a checkpoint file (to start finetuning for another task). Difference between `resume-from` and `load-from`: -`resume-from` loads both the model weights and optimizer status, and the iteration number is also inherited from the specified checkpoint. It is usually used for resuming the training process that is interrupted accidentally. -`load-from` only loads the model weights and the training iteration starts from 0. It is usually used for finetuning. +- `resume-from` loads both the model weights and optimizer state including the iteration number. +- `load-from` loads only the model weights, starts the training from iteration 0. ### Train with multiple machines diff --git a/tools/train.py b/tools/train.py index 26e8274..34c097f 100644 --- a/tools/train.py +++ b/tools/train.py @@ -20,6 +20,8 @@ def parse_args(): parser = argparse.ArgumentParser(description='Train a segmentor') parser.add_argument('config', help='train config file path') parser.add_argument('--work_dir', help='the dir to save logs and models') + parser.add_argument( + '--load-from', help='the checkpoint file to load weights from') parser.add_argument( '--resume-from', help='the checkpoint file to resume from') parser.add_argument( @@ -76,6 +78,8 @@ def main(): # use config filename as default work_dir if cfg.work_dir is None cfg.work_dir = osp.join('./work_dirs', osp.splitext(osp.basename(args.config))[0]) + if args.load_from is not None: + cfg.load_from = args.load_from if args.resume_from is not None: cfg.resume_from = args.resume_from if args.gpu_ids is not None: