From febc0365de294ab08fcfa1768b929473f6e47426 Mon Sep 17 00:00:00 2001 From: "yuze.zyz" Date: Sat, 13 May 2023 12:12:04 +0800 Subject: [PATCH] Support FlexTrain and update the structure of trainer 1. Refactor training_args 2. Refactor hooks 3. Add train_id for push_to_hub 4. Support both output_dir/output_sub_dir for checkpoint_hooks 5. Support copy when hardlink fails when checkpointing 6. Support mixed dataset config file as a CLI argument 7. Add eval txt in output folder Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12384253 * support the ignorance of file pattern --- .../finetune_image_classification.py | 2 +- .../finetune_multi_modal_embedding.py | 4 +- .../finetune_text_classification.py | 57 +- .../pytorch/text_classification/run_train.sh | 5 +- .../finetune_text_generation.py | 48 +- .../pytorch/text_generation/run_train_gpt3.sh | 5 +- .../finetune_token_classification.py | 4 +- modelscope/__init__.py | 79 +- modelscope/hub/api.py | 19 +- modelscope/hub/file_download.py | 1 - modelscope/hub/push_to_hub.py | 44 +- modelscope/models/audio/tts/voice.py | 4 +- modelscope/pipelines/__init__.py | 2 - modelscope/preprocessors/__init__.py | 2 +- modelscope/trainers/__init__.py | 6 +- modelscope/trainers/cli_argument_parser.py | 151 ++ modelscope/trainers/default_config.py | 37 +- modelscope/trainers/hooks/__init__.py | 6 +- .../trainers/hooks/checkpoint/__init__.py | 2 + .../hooks/checkpoint/checkpoint_hook.py | 435 ++++++ .../hooks/checkpoint/checkpoint_processor.py | 276 ++++ .../hooks/checkpoint/load_checkpoint_hook.py | 138 ++ modelscope/trainers/hooks/checkpoint_hook.py | 749 ---------- .../hooks/compression/sparsity_hook.py | 1 - .../trainers/hooks/distributed/__init__.py | 3 + .../hooks/{ => distributed}/ddp_hook.py | 6 +- .../hooks/{ => distributed}/deepspeed_hook.py | 145 +- .../hooks/{ => distributed}/megatron_hook.py | 239 +-- modelscope/trainers/hooks/early_stop_hook.py | 30 +- modelscope/trainers/hooks/evaluation_hook.py | 30 +- modelscope/trainers/hooks/hook.py | 72 +- .../trainers/hooks/lr_scheduler_hook.py | 132 +- .../hooks/optimizer/apex_optimizer_hook.py | 72 +- modelscope/trainers/hooks/optimizer/base.py | 89 +- .../hooks/optimizer/torch_optimizer_hook.py | 84 +- modelscope/trainers/trainer.py | 72 +- modelscope/trainers/training_args.py | 1315 +++++++---------- modelscope/utils/ast_utils.py | 6 +- modelscope/utils/checkpoint.py | 3 +- tests/hub/test_hub_upload.py | 1 + tests/trainers/cli/__init__.py | 0 tests/trainers/cli/test_cli.py | 52 + .../trainers/hooks/test_lr_scheduler_hook.py | 3 + tests/trainers/hooks/test_optimizer_hook.py | 1 + tests/trainers/test_trainer_with_nlp.py | 53 +- tests/trainers/test_training_args.py | 10 +- 46 files changed, 2394 insertions(+), 2101 deletions(-) create mode 100644 modelscope/trainers/cli_argument_parser.py create mode 100644 modelscope/trainers/hooks/checkpoint/__init__.py create mode 100644 modelscope/trainers/hooks/checkpoint/checkpoint_hook.py create mode 100644 modelscope/trainers/hooks/checkpoint/checkpoint_processor.py create mode 100644 modelscope/trainers/hooks/checkpoint/load_checkpoint_hook.py delete mode 100644 modelscope/trainers/hooks/checkpoint_hook.py create mode 100644 modelscope/trainers/hooks/distributed/__init__.py rename modelscope/trainers/hooks/{ => distributed}/ddp_hook.py (89%) rename modelscope/trainers/hooks/{ => distributed}/deepspeed_hook.py (64%) rename modelscope/trainers/hooks/{ => distributed}/megatron_hook.py (70%) create mode 100644 tests/trainers/cli/__init__.py create mode 100644 tests/trainers/cli/test_cli.py diff --git a/examples/pytorch/image_classification/finetune_image_classification.py b/examples/pytorch/image_classification/finetune_image_classification.py index 4e96c2cd..1e7fceb1 100644 --- a/examples/pytorch/image_classification/finetune_image_classification.py +++ b/examples/pytorch/image_classification/finetune_image_classification.py @@ -3,8 +3,8 @@ from dataclasses import dataclass, field from modelscope.metainfo import Trainers from modelscope.msdatasets.ms_dataset import MsDataset +from modelscope.trainers.args import TrainingArgs from modelscope.trainers.builder import build_trainer -from modelscope.trainers.training_args import TrainingArgs @dataclass diff --git a/examples/pytorch/multi_modal_embedding/finetune_multi_modal_embedding.py b/examples/pytorch/multi_modal_embedding/finetune_multi_modal_embedding.py index cc7da842..b3c325cf 100644 --- a/examples/pytorch/multi_modal_embedding/finetune_multi_modal_embedding.py +++ b/examples/pytorch/multi_modal_embedding/finetune_multi_modal_embedding.py @@ -5,8 +5,8 @@ from functools import partial from modelscope.metainfo import Trainers from modelscope.msdatasets import MsDataset from modelscope.trainers import build_trainer -from modelscope.trainers.training_args import (TrainingArgs, get_flatten_value, - set_flatten_value) +from modelscope.trainers.args import (TrainingArgs, get_flatten_value, + set_flatten_value) @dataclass diff --git a/examples/pytorch/text_classification/finetune_text_classification.py b/examples/pytorch/text_classification/finetune_text_classification.py index 7747bc25..73d15783 100644 --- a/examples/pytorch/text_classification/finetune_text_classification.py +++ b/examples/pytorch/text_classification/finetune_text_classification.py @@ -1,26 +1,17 @@ import os from dataclasses import dataclass, field -from modelscope.msdatasets import MsDataset -from modelscope.trainers import EpochBasedTrainer, build_trainer -from modelscope.trainers.training_args import TrainingArgs +from modelscope import (EpochBasedTrainer, MsDataset, TrainingArgs, + build_dataset_from_file, build_trainer) -def get_labels(cfg, metadata): - label2id = cfg.safe_get(metadata['cfg_node']) - if label2id is not None: - return ','.join(label2id.keys()) - - -def set_labels(cfg, labels, metadata): +def set_labels(labels): if isinstance(labels, str): labels = labels.split(',') - cfg.merge_from_dict( - {metadata['cfg_node']: {label: id - for id, label in enumerate(labels)}}) + return {label: id for id, label in enumerate(labels)} -@dataclass +@dataclass(init=False) class TextClassificationArguments(TrainingArgs): first_sequence: str = field( @@ -49,7 +40,6 @@ class TextClassificationArguments(TrainingArgs): metadata={ 'help': 'The labels of the dataset', 'cfg_node': 'preprocessor.label2id', - 'cfg_getter': get_labels, 'cfg_setter': set_labels, }) @@ -60,30 +50,39 @@ class TextClassificationArguments(TrainingArgs): 'cfg_node': 'preprocessor.type' }) - def __call__(self, config): - config = super().__call__(config) - config.model['num_labels'] = len(self.labels) - if config.train.lr_scheduler.type == 'LinearLR': - config.train.lr_scheduler['total_iters'] = \ - int(len(train_dataset) / self.per_device_train_batch_size) * self.max_epochs - return config + +config, args = TextClassificationArguments().parse_cli().to_config() + +print(config, args) -args = TextClassificationArguments.from_cli( - task='text-classification', eval_metrics='seq-cls-metric') +def cfg_modify_fn(cfg): + if args.use_model_config: + cfg.merge_from_dict(config) + else: + cfg = config + cfg.model['num_labels'] = len(cfg.preprocessor.label2id) + if cfg.train.lr_scheduler.type == 'LinearLR': + cfg.train.lr_scheduler['total_iters'] = \ + int(len(train_dataset) / cfg.train.dataloader.batch_size_per_gpu) * cfg.train.max_epochs + return cfg -print(args) -dataset = MsDataset.load(args.dataset_name, subset_name=args.subset_name) -train_dataset = dataset['train'] -validation_dataset = dataset['validation'] +if args.dataset_json_file is None: + dataset = MsDataset.load( + args.train_dataset_name, subset_name=args.train_subset_name) + train_dataset = dataset['train'] + validation_dataset = dataset['validation'] +else: + train_dataset, validation_dataset = build_dataset_from_file( + args.dataset_json_file) kwargs = dict( model=args.model, train_dataset=train_dataset, eval_dataset=validation_dataset, seed=args.seed, - cfg_modify_fn=args) + cfg_modify_fn=cfg_modify_fn) os.environ['LOCAL_RANK'] = str(args.local_rank) trainer: EpochBasedTrainer = build_trainer(name='trainer', default_args=kwargs) diff --git a/examples/pytorch/text_classification/run_train.sh b/examples/pytorch/text_classification/run_train.sh index 93c23d0d..f46fb72a 100644 --- a/examples/pytorch/text_classification/run_train.sh +++ b/examples/pytorch/text_classification/run_train.sh @@ -1,7 +1,7 @@ PYTHONPATH=. python examples/pytorch/text_classification/finetune_text_classification.py \ --model 'damo/nlp_structbert_backbone_base_std' \ - --dataset_name 'clue' \ - --subset_name 'tnews' \ + --train_dataset_name 'clue' \ + --train_subset_name 'tnews' \ --first_sequence 'sentence' \ --preprocessor.label label \ --model.num_labels 15 \ @@ -10,3 +10,4 @@ PYTHONPATH=. python examples/pytorch/text_classification/finetune_text_classific --train.dataloader.workers_per_gpu 0 \ --evaluation.dataloader.workers_per_gpu 0 \ --train.optimizer.lr 1e-5 \ + --use_model_config true \ diff --git a/examples/pytorch/text_generation/finetune_text_generation.py b/examples/pytorch/text_generation/finetune_text_generation.py index 7a140a0c..f9e47bf3 100644 --- a/examples/pytorch/text_generation/finetune_text_generation.py +++ b/examples/pytorch/text_generation/finetune_text_generation.py @@ -1,12 +1,11 @@ from dataclasses import dataclass, field +from modelscope import (EpochBasedTrainer, MsDataset, TrainingArgs, + build_trainer) from modelscope.metainfo import Trainers -from modelscope.msdatasets import MsDataset -from modelscope.trainers import EpochBasedTrainer, build_trainer -from modelscope.trainers.training_args import TrainingArgs -@dataclass +@dataclass(init=False) class TextGenerationArguments(TrainingArgs): trainer: str = field( @@ -67,30 +66,35 @@ class TextGenerationArguments(TrainingArgs): 'help': 'Whether to use MegatronHook', }) - def __call__(self, config): - config = super().__call__(config) - if config.train.lr_scheduler.type == 'noam': - config.train.lr_scheduler = { - 'type': 'LambdaLR', - 'lr_lambda': noam_lambda, - 'options': { - 'by_epoch': False - } - } - if self.use_megatron: - config.train.hooks.append({'type': 'MegatronHook'}) - return config - def noam_lambda(current_step: int): current_step += 1 return min(current_step**(-0.5), current_step * 100**(-1.5)) -args = TextGenerationArguments.from_cli(task='text-generation') -print(args) +config, args = TextGenerationArguments().parse_cli().to_config() +print(config, args) -dataset = MsDataset.load(args.dataset_name) + +def cfg_modify_fn(cfg): + if args.use_model_config: + cfg.merge_from_dict(config) + else: + cfg = config + if cfg.train.lr_scheduler.type == 'noam': + cfg.train.lr_scheduler = { + 'type': 'LambdaLR', + 'lr_lambda': noam_lambda, + 'options': { + 'by_epoch': False + } + } + if args.use_megatron: + cfg.train.hooks.append({'type': 'MegatronHook'}) + return cfg + + +dataset = MsDataset.load(args.train_dataset_name) train_dataset = dataset['train'] eval_dataset = dataset['validation' if 'validation' in dataset else 'test'] @@ -100,7 +104,7 @@ kwargs = dict( eval_dataset=eval_dataset, seed=args.seed, work_dir=args.work_dir, - cfg_modify_fn=args) + cfg_modify_fn=cfg_modify_fn) trainer: EpochBasedTrainer = build_trainer( name=args.trainer, default_args=kwargs) diff --git a/examples/pytorch/text_generation/run_train_gpt3.sh b/examples/pytorch/text_generation/run_train_gpt3.sh index a20a5bb2..fd37b42c 100644 --- a/examples/pytorch/text_generation/run_train_gpt3.sh +++ b/examples/pytorch/text_generation/run_train_gpt3.sh @@ -8,7 +8,7 @@ PYTHONPATH=. torchrun --nproc_per_node $WORLD_SIZE examples/pytorch/text_generat --trainer 'nlp-gpt3-trainer' \ --work_dir './tmp' \ --model 'damo/nlp_gpt3_text-generation_1.3B' \ - --dataset_name 'chinese-poetry-collection' \ + --train_dataset_name 'chinese-poetry-collection' \ --preprocessor 'text-gen-jieba-tokenizer' \ --src_txt 'text1' \ --tgt_txt 'text2' \ @@ -20,4 +20,5 @@ PYTHONPATH=. torchrun --nproc_per_node $WORLD_SIZE examples/pytorch/text_generat --world_size $WORLD_SIZE \ --tensor_model_parallel_size $TENSOR_MODEL_PARALLEL_SIZE \ --use_megatron true \ - # --dataset_name 'DuReader_robust-QG' \ # input&output + --use_model_config true \ + # --train_dataset_name 'DuReader_robust-QG' \ # input&output diff --git a/examples/pytorch/token_classification/finetune_token_classification.py b/examples/pytorch/token_classification/finetune_token_classification.py index cf51ed22..e06f6f48 100644 --- a/examples/pytorch/token_classification/finetune_token_classification.py +++ b/examples/pytorch/token_classification/finetune_token_classification.py @@ -3,8 +3,8 @@ from dataclasses import dataclass, field from modelscope.metainfo import Trainers from modelscope.msdatasets import MsDataset from modelscope.trainers import build_trainer -from modelscope.trainers.training_args import (TrainingArgs, get_flatten_value, - set_flatten_value) +from modelscope.trainers.args import (TrainingArgs, get_flatten_value, + set_flatten_value) @dataclass diff --git a/modelscope/__init__.py b/modelscope/__init__.py index 81fdf505..f7553958 100644 --- a/modelscope/__init__.py +++ b/modelscope/__init__.py @@ -1,4 +1,79 @@ # Copyright (c) Alibaba, Inc. and its affiliates. -from .version import __release_datetime__, __version__ +from typing import TYPE_CHECKING -__all__ = ['__version__', '__release_datetime__'] +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .version import __release_datetime__, __version__ + from .trainers import EpochBasedTrainer, TrainingArgs, build_dataset_from_file + from .trainers import Hook, Priority + from .exporters import Exporter + from .exporters import TfModelExporter + from .exporters import TorchModelExporter + from .hub.api import HubApi + from .hub.snapshot_download import snapshot_download + from .hub.push_to_hub import push_to_hub, push_to_hub_async + from .hub.check_model import check_model_is_id, check_local_model_is_latest + from .metrics import AudioNoiseMetric, Metric, task_default_metrics, ImageColorEnhanceMetric, ImageDenoiseMetric, \ + ImageInstanceSegmentationCOCOMetric, ImagePortraitEnhancementMetric, SequenceClassificationMetric, \ + TextGenerationMetric, TokenClassificationMetric, VideoSummarizationMetric, MovieSceneSegmentationMetric, \ + AccuracyMetric, BleuMetric, ImageInpaintingMetric, ReferringVideoObjectSegmentationMetric, \ + VideoFrameInterpolationMetric, VideoStabilizationMetric, VideoSuperResolutionMetric, PplMetric, \ + ImageQualityAssessmentDegradationMetric, ImageQualityAssessmentMosMetric, TextRankingMetric, \ + LossMetric, ImageColorizationMetric, OCRRecognitionMetric + from .models import Model, TorchModel + from .preprocessors import Preprocessor + from .pipelines import Pipeline, pipeline + from .utils.hub import read_config, create_model_if_not_exist + from .utils.logger import get_logger + from .msdatasets import MsDataset + +else: + _import_structure = { + 'version': ['__release_datetime__', '__version__'], + 'trainers': [ + 'EpochBasedTrainer', 'TrainingArgs', 'Hook', 'Priority', + 'build_dataset_from_file' + ], + 'exporters': [ + 'Exporter', + 'TfModelExporter', + 'TorchModelExporter', + ], + 'hub.api': ['HubApi'], + 'hub.snapshot_download': ['snapshot_download'], + 'hub.push_to_hub': ['push_to_hub', 'push_to_hub_async'], + 'hub.check_model': + ['check_model_is_id', 'check_local_model_is_latest'], + 'metrics': [ + 'AudioNoiseMetric', 'Metric', 'task_default_metrics', + 'ImageColorEnhanceMetric', 'ImageDenoiseMetric', + 'ImageInstanceSegmentationCOCOMetric', + 'ImagePortraitEnhancementMetric', 'SequenceClassificationMetric', + 'TextGenerationMetric', 'TokenClassificationMetric', + 'VideoSummarizationMetric', 'MovieSceneSegmentationMetric', + 'AccuracyMetric', 'BleuMetric', 'ImageInpaintingMetric', + 'ReferringVideoObjectSegmentationMetric', + 'VideoFrameInterpolationMetric', 'VideoStabilizationMetric', + 'VideoSuperResolutionMetric', 'PplMetric', + 'ImageQualityAssessmentDegradationMetric', + 'ImageQualityAssessmentMosMetric', 'TextRankingMetric', + 'LossMetric', 'ImageColorizationMetric', 'OCRRecognitionMetric' + ], + 'models': ['Model', 'TorchModel'], + 'preprocessors': ['Preprocessor'], + 'pipelines': ['Pipeline', 'pipeline'], + 'utils.hub': ['read_config', 'create_model_if_not_exist'], + 'utils.logger': ['get_logger'], + 'msdatasets': ['MsDataset'] + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index 0991dbb9..e3436aea 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -6,6 +6,7 @@ import functools import os import pickle import platform +import re import shutil import tempfile import uuid @@ -19,7 +20,6 @@ import requests from requests import Session from requests.adapters import HTTPAdapter, Retry -from modelscope import __version__ from modelscope.hub.constants import (API_HTTP_CLIENT_TIMEOUT, API_RESPONSE_FIELD_DATA, API_RESPONSE_FIELD_EMAIL, @@ -161,6 +161,7 @@ class HubApi: 'Visibility': visibility, # server check 'License': license, 'OriginalModelId': original_model_id, + 'TrainId': os.environ.get('MODELSCOPE_TRAIN_ID', ''), } r = self.session.post( path, json=body, cookies=cookies, headers=self.headers) @@ -237,8 +238,10 @@ class HubApi: license: Optional[str] = Licenses.APACHE_V2, chinese_name: Optional[str] = None, commit_message: Optional[str] = 'upload model', + tag: Optional[str] = None, revision: Optional[str] = DEFAULT_REPOSITORY_REVISION, - original_model_id: Optional[str] = None): + original_model_id: Optional[str] = None, + ignore_file_pattern: Optional[Union[List[str], str]] = None): """Upload model from a given directory to given repository. A valid model directory must contain a configuration.json file. @@ -269,10 +272,13 @@ class HubApi: chinese name of the new created model. commit_message(`str`, *optional*, defaults to `None`): commit message of the push request. + tag(`str`, *optional*, defaults to `None`): + The tag on this commit revision (`str`, *optional*, default to DEFAULT_MODEL_REVISION): which branch to push. If the branch is not exists, It will create a new branch and push to it. original_model_id (str, optional): The base model id which this model is trained from + ignore_file_pattern (`Union[List[str], str]`, optional): The file pattern to ignore uploading Raises: InvalidParameter: Parameter invalid. @@ -293,6 +299,10 @@ class HubApi: if cookies is None: raise NotLoginException('Must login before upload!') files_to_save = os.listdir(model_dir) + if ignore_file_pattern is None: + ignore_file_pattern = [] + if isinstance(ignore_file_pattern, str): + ignore_file_pattern = [ignore_file_pattern] try: self.get_model(model_id=model_id) except Exception: @@ -326,6 +336,8 @@ class HubApi: shutil.rmtree(src, ignore_errors=True) for f in files_to_save: if f[0] != '.': + if any([re.search(pattern, f) is not None for pattern in ignore_file_pattern]): + continue src = os.path.join(model_dir, f) if os.path.isdir(src): shutil.copytree(src, os.path.join(tmp_dir, f)) @@ -339,6 +351,8 @@ class HubApi: commit_message=commit_message, local_branch=revision, remote_branch=revision) + if tag is not None: + repo.tag_and_push(tag, tag) except Exception: raise finally: @@ -928,6 +942,7 @@ class ModelScopeConfig: if MODELSCOPE_CLOUD_USERNAME in os.environ: user_name = os.environ[MODELSCOPE_CLOUD_USERNAME] + from modelscope import __version__ ua = 'modelscope/%s; python/%s; session_id/%s; platform/%s; processor/%s; env/%s; user/%s' % ( __version__, platform.python_version(), diff --git a/modelscope/hub/file_download.py b/modelscope/hub/file_download.py index 380d2432..6d3ad63d 100644 --- a/modelscope/hub/file_download.py +++ b/modelscope/hub/file_download.py @@ -12,7 +12,6 @@ import requests from requests.adapters import Retry from tqdm import tqdm -from modelscope import __version__ from modelscope.hub.api import HubApi, ModelScopeConfig from modelscope.hub.constants import (API_FILE_DOWNLOAD_CHUNK_SIZE, API_FILE_DOWNLOAD_RETRY_TIMES, diff --git a/modelscope/hub/push_to_hub.py b/modelscope/hub/push_to_hub.py index ee7b240e..d117cc7f 100644 --- a/modelscope/hub/push_to_hub.py +++ b/modelscope/hub/push_to_hub.py @@ -4,8 +4,8 @@ import concurrent.futures import os from modelscope.hub.api import HubApi -from modelscope.hub.constants import Licenses, ModelVisibility -from modelscope.hub.errors import NotExistError +from modelscope.hub.constants import ModelVisibility +from modelscope.utils.constant import DEFAULT_REPOSITORY_REVISION from modelscope.utils.logger import get_logger logger = get_logger() @@ -18,7 +18,10 @@ def _api_push_to_hub(repo_name, token, private=True, commit_message='', - source_repo=''): + tag=None, + source_repo='', + ignore_file_pattern=None, + revision=DEFAULT_REPOSITORY_REVISION): try: api = HubApi() api.login(token) @@ -29,7 +32,10 @@ def _api_push_to_hub(repo_name, if not private else ModelVisibility.PRIVATE, chinese_name=repo_name, commit_message=commit_message, - original_model_id=source_repo) + tag=tag, + original_model_id=source_repo, + ignore_file_pattern=ignore_file_pattern, + revision=revision) commit_message = commit_message or 'No commit message' logger.info( f'Successfully upload the model to {repo_name} with message: {commit_message}' @@ -48,7 +54,10 @@ def push_to_hub(repo_name, private=True, retry=3, commit_message='', - source_repo=''): + tag=None, + source_repo='', + ignore_file_pattern=None, + revision=DEFAULT_REPOSITORY_REVISION): """ Args: repo_name: The repo name for the modelhub repo @@ -57,13 +66,18 @@ def push_to_hub(repo_name, private: If is a private repo, default True retry: Retry times if something error in uploading, default 3 commit_message: The commit message + tag: The tag of this commit source_repo: The source repo (model id) which this model comes from - + ignore_file_pattern: The file pattern to be ignored in uploading. + revision: The branch to commit to Returns: The boolean value to represent whether the model is uploaded. """ if token is None: token = os.environ.get('MODELSCOPE_API_TOKEN') + if ignore_file_pattern is None: + ignore_file_pattern = os.environ.get('UPLOAD_IGNORE_FILE_PATTERN') + assert repo_name is not None assert token is not None, 'Either pass in a token or to set `MODELSCOPE_API_TOKEN` in the environment variables.' assert os.path.isdir(output_dir) assert 'configuration.json' in os.listdir(output_dir) or 'configuration.yaml' in os.listdir(output_dir) \ @@ -73,7 +87,8 @@ def push_to_hub(repo_name, f'Uploading {output_dir} to {repo_name} with message {commit_message}') for i in range(retry): if _api_push_to_hub(repo_name, output_dir, token, private, - commit_message, source_repo): + commit_message, tag, source_repo, + ignore_file_pattern, revision): return True return False @@ -83,7 +98,10 @@ def push_to_hub_async(repo_name, token=None, private=True, commit_message='', - source_repo=''): + tag=None, + source_repo='', + ignore_file_pattern=None, + revision=DEFAULT_REPOSITORY_REVISION): """ Args: repo_name: The repo name for the modelhub repo @@ -91,13 +109,18 @@ def push_to_hub_async(repo_name, token: The user api token, function will check the `MODELSCOPE_API_TOKEN` variable if this argument is None private: If is a private repo, default True commit_message: The commit message + tag: The tag of this commit source_repo: The source repo (model id) which this model comes from - + ignore_file_pattern: The file pattern to be ignored in uploading + revision: The branch to commit to Returns: A handler to check the result and the status """ if token is None: token = os.environ.get('MODELSCOPE_API_TOKEN') + if ignore_file_pattern is None: + ignore_file_pattern = os.environ.get('UPLOAD_IGNORE_FILE_PATTERN') + assert repo_name is not None assert token is not None, 'Either pass in a token or to set `MODELSCOPE_API_TOKEN` in the environment variables.' assert os.path.isdir(output_dir) assert 'configuration.json' in os.listdir(output_dir) or 'configuration.yaml' in os.listdir(output_dir) \ @@ -106,4 +129,5 @@ def push_to_hub_async(repo_name, logger.info( f'Uploading {output_dir} to {repo_name} with message {commit_message}') return _executor.submit(_api_push_to_hub, repo_name, output_dir, token, - private, commit_message, source_repo) + private, commit_message, tag, source_repo, + ignore_file_pattern, revision) diff --git a/modelscope/models/audio/tts/voice.py b/modelscope/models/audio/tts/voice.py index 645a528f..ed9edf43 100644 --- a/modelscope/models/audio/tts/voice.py +++ b/modelscope/models/audio/tts/voice.py @@ -17,11 +17,9 @@ from kantts.train.trainer import GAN_Trainer, Sambert_Trainer, distributed_init from kantts.utils.ling_unit.ling_unit import KanTtsLinguisticUnit from torch.utils.data import DataLoader -from modelscope import __version__ from modelscope.utils.audio.audio_utils import TtsCustomParams from modelscope.utils.audio.tts_exceptions import ( TtsModelConfigurationException, TtsModelNotExistsException) -from modelscope.utils.constant import ModelFile, Tasks from modelscope.utils.logger import get_logger logger = get_logger() @@ -394,6 +392,7 @@ class Voice: logger.info(f'TRAINING steps: {train_max_steps}') config['create_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()) + from modelscope import __version__ config['modelscope_version'] = __version__ with open(os.path.join(stage_dir, 'config.yaml'), 'w') as f: @@ -558,6 +557,7 @@ class Voice: logger.info(f'resume from: {resume_from}') config['create_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()) + from modelscope import __version__ config['modelscope_version'] = __version__ with open(os.path.join(stage_dir, 'config.yaml'), 'w') as f: diff --git a/modelscope/pipelines/__init__.py b/modelscope/pipelines/__init__.py index 71fe307b..d98a7af9 100644 --- a/modelscope/pipelines/__init__.py +++ b/modelscope/pipelines/__init__.py @@ -1,7 +1,5 @@ # Copyright (c) Alibaba, Inc. and its affiliates. -from typing import TYPE_CHECKING -from modelscope.utils.import_utils import LazyImportModule from . import audio, cv, multi_modal, nlp from .base import Pipeline from .builder import pipeline diff --git a/modelscope/preprocessors/__init__.py b/modelscope/preprocessors/__init__.py index fab055db..3bbca124 100644 --- a/modelscope/preprocessors/__init__.py +++ b/modelscope/preprocessors/__init__.py @@ -35,7 +35,7 @@ if TYPE_CHECKING: TextErrorCorrectionPreprocessor, TextGenerationT5Preprocessor, WordAlignmentPreprocessor, TextGenerationTransformersPreprocessor, Tokenize, WordSegmentationBlankSetToLabelPreprocessor, - CodeGeeXPreprocessor, MGLMSummarizationPreprocessor, + MGLMSummarizationPreprocessor, ZeroShotClassificationTransformersPreprocessor, TextGenerationJiebaPreprocessor, SentencePiecePreprocessor, DialogIntentPredictionPreprocessor, DialogModelingPreprocessor, diff --git a/modelscope/trainers/__init__.py b/modelscope/trainers/__init__.py index 90f73a7f..0d20fe00 100644 --- a/modelscope/trainers/__init__.py +++ b/modelscope/trainers/__init__.py @@ -15,6 +15,8 @@ if TYPE_CHECKING: from .nlp import SequenceClassificationTrainer, TextRankingTrainer, SiameseUIETrainer from .nlp_trainer import NlpEpochBasedTrainer, VecoTrainer from .trainer import EpochBasedTrainer + from .training_args import TrainingArgs, build_dataset_from_file + from .hooks import Hook, Priority else: _import_structure = { @@ -32,7 +34,9 @@ else: 'SiameseUIETrainer' ], 'nlp_trainer': ['NlpEpochBasedTrainer', 'VecoTrainer'], - 'trainer': ['EpochBasedTrainer'] + 'trainer': ['EpochBasedTrainer'], + 'training_args': ['TrainingArgs', 'build_dataset_from_file'], + 'hooks': ['Hook'] } import sys diff --git a/modelscope/trainers/cli_argument_parser.py b/modelscope/trainers/cli_argument_parser.py new file mode 100644 index 00000000..f183b9ea --- /dev/null +++ b/modelscope/trainers/cli_argument_parser.py @@ -0,0 +1,151 @@ +from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentParser +from dataclasses import fields +from typing import List + + +class CliArgumentParser(ArgumentParser): + """ Argument Parser to define and parse command-line args for training. + + Args: + training_args: dict or list of dict which defines different + paramters for training. + """ + + def __init__(self, training_args=None, **kwargs): + if 'formatter_class' not in kwargs: + kwargs['formatter_class'] = ArgumentDefaultsHelpFormatter + super().__init__(**kwargs) + self.training_args = training_args + self.define_args() + + def get_manual_args(self, args): + return [arg[2:] for arg in args if arg.startswith('--')] + + def _parse_known_args(self, args: List = None, namespace=None): + self.model_id = namespace.model if namespace is not None else None + if '--model' in args: + self.model_id = args[args.index('--model') + 1] + self.manual_args = self.get_manual_args(args) + return super()._parse_known_args(args, namespace) + + def print_help(self, file=None): + return super().print_help(file) + + def define_args(self): + if self.training_args is not None: + for f in fields(self.training_args): + arg_name = f.name + arg_attr = getattr(self.training_args, f.name) + name = f'--{arg_name}' + kwargs = dict(type=f.type, help=f.metadata['help']) + kwargs['default'] = arg_attr + + if 'choices' in f.metadata: + kwargs['choices'] = f.metadata['choices'] + + kwargs['action'] = SingleAction + self.add_argument(name, **kwargs) + + +class DictAction(Action): + """ + argparse action to split an argument into KEY=VALUE form + on the first = and append to a dictionary. List options can + be passed as comma separated values, i.e 'KEY=V1,V2,V3', or with explicit + brackets, i.e. 'KEY=[V1,V2,V3]'. It also support nested brackets to build + list/tuple values. e.g. 'KEY=[(V1,V2),(V3,V4)]' + """ + + @staticmethod + def parse_int_float_bool_str(val): + try: + return int(val) + except ValueError: + pass + try: + return float(val) + except ValueError: + pass + if val.lower() in ['true', 'false']: + return val.lower() == 'true' + if val == 'None': + return None + return val + + @staticmethod + def parse_iterable(val): + """Parse iterable values in the string. + All elements inside '()' or '[]' are treated as iterable values. + Args: + val (str): Value string. + Returns: + list | tuple: The expanded list or tuple from the string. + Examples: + >>> DictAction._parse_iterable('1,2,3') + [1, 2, 3] + >>> DictAction._parse_iterable('[a, b, c]') + ['a', 'b', 'c'] + >>> DictAction._parse_iterable('[(1, 2, 3), [a, b], c]') + [(1, 2, 3), ['a', 'b'], 'c'] + """ + + def find_next_comma(string): + """Find the position of next comma in the string. + If no ',' is found in the string, return the string length. All + chars inside '()' and '[]' are treated as one element and thus ',' + inside these brackets are ignored. + """ + assert (string.count('(') == string.count(')')) and ( + string.count('[') + == string.count(']')), f'Imbalanced brackets exist in {string}' + end = len(string) + for idx, char in enumerate(string): + pre = string[:idx] + # The string before this ',' is balanced + if ((char == ',') and (pre.count('(') == pre.count(')')) + and (pre.count('[') == pre.count(']'))): + end = idx + break + return end + + # Strip ' and " characters and replace whitespace. + val = val.strip('\'\"').replace(' ', '') + is_tuple = False + if val.startswith('(') and val.endswith(')'): + is_tuple = True + val = val[1:-1] + elif val.startswith('[') and val.endswith(']'): + val = val[1:-1] + elif ',' not in val: + # val is a single value + return DictAction.parse_int_float_bool_str(val) + + values = [] + while len(val) > 0: + comma_idx = find_next_comma(val) + element = DictAction.parse_iterable(val[:comma_idx]) + values.append(element) + val = val[comma_idx + 1:] + if is_tuple: + values = tuple(values) + return values + + def __call__(self, parser, namespace, values, option_string): + options = {} + for kv in values: + key, val = kv.split('=', maxsplit=1) + options[key] = self.parse_iterable(val) + setattr(namespace, self.dest, options) + + +class SingleAction(DictAction): + """ Argparse action to convert value to tuple or list or nested structure of + list and tuple, i.e 'V1,V2,V3', or with explicit brackets, i.e. '[V1,V2,V3]'. + It also support nested brackets to build list/tuple values. e.g. '[(V1,V2),(V3,V4)]' + """ + + def __call__(self, parser, namespace, value, option_string): + if isinstance(value, str): + setattr(namespace, self.dest, self.parse_iterable(value)) + else: + setattr(namespace, self.dest, value) diff --git a/modelscope/trainers/default_config.py b/modelscope/trainers/default_config.py index 51a0df40..bb272695 100644 --- a/modelscope/trainers/default_config.py +++ b/modelscope/trainers/default_config.py @@ -4,38 +4,6 @@ from typing import Dict, List, Optional, Tuple from modelscope.utils.config import Config -DEFAULT_CONFIG = Config({ - 'framework': 'pytorch', - 'train': { - 'work_dir': '/tmp', - 'max_epochs': 10, - 'dataloader': { - 'batch_size_per_gpu': 16, - 'workers_per_gpu': 0 - }, - 'optimizer': { - 'type': 'SGD', - 'lr': 1e-3 - }, - 'lr_scheduler': { - 'type': 'StepLR', - 'step_size': 2 - }, - 'checkpoint': { - 'period': { - 'interval': 1 - } - } - }, - 'evaluation': { - 'dataloader': { - 'batch_size_per_gpu': 16, - 'workers_per_gpu': 0, - 'shuffle': False - }, - } -}) - DEFAULT_HOOKS_CONFIG = { 'train.hooks': [{ 'type': 'CheckpointHook', @@ -68,7 +36,7 @@ def merge_cfg(cfg: Config): def merge_hooks(cfg: Config) -> List[Dict]: - hooks = cfg.train.hooks.copy() + hooks = getattr(cfg.train, 'hooks', []).copy() for hook_type, key_chain in _HOOK_KEY_CHAIN_MAP.items(): hook = _key_chain_to_hook(cfg, key_chain, hook_type) if hook is not None: @@ -107,7 +75,8 @@ def _check_basic_hook(cfg: Config, key_chain: str, hook_type: str) -> bool: if cfg.safe_get(key_chain) is None: return False hooks = list( - filter(lambda hook: hook['type'] == hook_type, cfg.train.hooks)) + filter(lambda hook: hook['type'] == hook_type, + getattr(cfg.train, 'hooks', []))) assert len(hooks) == 0, f'The key_chain {key_chain} and the traditional hook ' \ f'cannot exist at the same time, ' \ f'please delete {hook_type} in the configuration file.' diff --git a/modelscope/trainers/hooks/__init__.py b/modelscope/trainers/hooks/__init__.py index 51677f25..5485f454 100644 --- a/modelscope/trainers/hooks/__init__.py +++ b/modelscope/trainers/hooks/__init__.py @@ -5,7 +5,6 @@ from modelscope.utils.import_utils import LazyImportModule if TYPE_CHECKING: from .builder import HOOKS, build_hook - from .checkpoint_hook import BestCkptSaverHook, CheckpointHook, LoadCheckpointHook from .early_stop_hook import EarlyStopHook from .compression import SparsityHook from .evaluation_hook import EvaluationHook @@ -16,6 +15,8 @@ if TYPE_CHECKING: from .optimizer import (ApexAMPOptimizerHook, NoneOptimizerHook, OptimizerHook, TorchAMPOptimizerHook) from .priority import Priority, get_priority + from .checkpoint import CheckpointHook, LoadCheckpointHook, BestCkptSaverHook + from .distributed import DDPHook, DeepspeedHook, MegatronHook else: _import_structure = { @@ -32,6 +33,9 @@ else: 'ApexAMPOptimizerHook', 'NoneOptimizerHook', 'OptimizerHook', 'TorchAMPOptimizerHook' ], + 'checkpoint': + ['CheckpointHook', 'LoadCheckpointHook', 'BestCkptSaverHook'], + 'distributed': ['DDPHook', 'DeepspeedHook', 'MegatronHook'], 'priority': ['Priority', 'get'] } diff --git a/modelscope/trainers/hooks/checkpoint/__init__.py b/modelscope/trainers/hooks/checkpoint/__init__.py new file mode 100644 index 00000000..e2abb272 --- /dev/null +++ b/modelscope/trainers/hooks/checkpoint/__init__.py @@ -0,0 +1,2 @@ +from .checkpoint_hook import BestCkptSaverHook, CheckpointHook +from .load_checkpoint_hook import LoadCheckpointHook diff --git a/modelscope/trainers/hooks/checkpoint/checkpoint_hook.py b/modelscope/trainers/hooks/checkpoint/checkpoint_hook.py new file mode 100644 index 00000000..e531a325 --- /dev/null +++ b/modelscope/trainers/hooks/checkpoint/checkpoint_hook.py @@ -0,0 +1,435 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +import random +import time +from typing import Optional + +import numpy as np +import torch + +from modelscope.hub.check_model import check_model_is_id +from modelscope.hub.push_to_hub import push_to_hub_async +from modelscope.metainfo import Hooks +from modelscope.trainers.hooks.builder import HOOKS +from modelscope.trainers.hooks.checkpoint.checkpoint_processor import \ + CheckpointProcessor +from modelscope.trainers.hooks.hook import Hook +from modelscope.trainers.hooks.priority import Priority +from modelscope.utils.constant import (DEFAULT_REPOSITORY_REVISION, LogKeys, + ModelFile) +from modelscope.utils.logger import get_logger +from modelscope.utils.torch_utils import is_master + + +class CheckpointStrategy: + by_epoch = 'by_epoch' + by_step = 'by_step' + no = 'no' + + +@HOOKS.register_module(module_name=Hooks.CheckpointHook) +class CheckpointHook(Hook): + """Save checkpoints periodically. + + Args: + save_strategy(str): The strategy to save checkpoint, can be `by_epoch`, `by_step` or `no` + interval (int): The frequency to save model. If `by_epoch=True`, + it means the number of epochs, else means the number of iterations + save_dir (str): The directory to save checkpoints. If is None, use `trainer.work_dir` + output_dir (str): The absolute path to save the output files for inference. If it's not specified, + the default dir is `{sub_dir}/output`. + save_last (bool): Whether to save the last checkpoint. Default: True. + max_checkpoint_num (int): The max number of checkpoint files, default None which means never delete anything. + If the number exceeding the limit, earlier checkpoints will be deleted first. + push_to_hub (bool): Whether push the checkpoint to modelhub. + hub_repo_id (str): The hub repo id. + hub_token (str): The token of the modelhub. You can also set the environment variable `MODELSCOPE_API_TOKEN`. + private_hub (bool): Whether push to a private hub, default True. + hub_revision (str): Which branch to push the model to, default is `master` + kwargs: + by_epoch (bool): Same with `save_strategy`, but has a higher priority, legacy argument. + output_sub_dir (str): The folder under the `save_dir` to save the output checkpoint for inference. + This argument is kept to fit the existing configs. + """ + + PRIORITY = Priority.LOW + + EVAL_RESULT_FILE = 'eval_result.txt' + + def __init__(self, + save_strategy: Optional[str] = CheckpointStrategy.by_epoch, + interval: Optional[int] = 0, + save_dir: Optional[str] = None, + output_dir: Optional[str] = None, + save_last: Optional[bool] = True, + max_checkpoint_num: Optional[int] = None, + push_to_hub: Optional[bool] = False, + hub_repo_id: Optional[str] = None, + hub_token: Optional[str] = None, + private_hub: Optional[bool] = True, + hub_revision: Optional[str] = DEFAULT_REPOSITORY_REVISION, + **kwargs): + self.interval = interval + self.save_dir = save_dir + if 'by_epoch' in kwargs: + self.save_strategy = CheckpointStrategy.by_epoch if kwargs[ + 'by_epoch'] else CheckpointStrategy.by_step + else: + self.save_strategy = save_strategy + if 'output_sub_dir' in kwargs: + self.output_sub_dir = kwargs['output_sub_dir'] + self.output_dir = None + else: + self.output_sub_dir = None + self.output_dir = output_dir + self.save_last = save_last + self.rng_state = None + self.push_to_hub = push_to_hub + self.hub_repo_id = hub_repo_id + self.hub_token = hub_token + self.private_hub = private_hub + self.hub_revision = hub_revision + self.tag = -1 + self.is_model_id = None + self.push_to_hub_future = None + self.max_checkpoint_num = None + if max_checkpoint_num is not None: + self.max_checkpoint_num = max(int(max_checkpoint_num), 1) + self.history_checkpoints = [] + self.processor = CheckpointProcessor() + + def set_processor(self, processor): + """ + The checkpoint hook accepts a processor to finish the actual saving/deleting action. + """ + self.processor = processor + + def before_run(self, trainer): + self.tag = -1 + if not self.save_dir: + self.save_dir = trainer.work_dir + if not self.output_dir: + if self.output_sub_dir: + self.output_dir = os.path.join(self.save_dir, + self.output_sub_dir) + else: + self.output_dir = os.path.join(self.save_dir, + ModelFile.TRAIN_OUTPUT_DIR) + + if not os.path.exists(self.save_dir): + os.makedirs(self.save_dir, exist_ok=True) + + if not hasattr(trainer, 'logger'): + self.logger = get_logger() + else: + self.logger = trainer.logger + + if is_master(): + output_dir = self.output_dir + # only global master prepares the output folder + self.processor.prepare_output(trainer, output_dir) + self.logger.info(f'Checkpoints will be saved to {self.save_dir}') + + def generate_prefix(self, trainer, save_strategy): + if save_strategy == CheckpointStrategy.by_epoch: + return f'{LogKeys.EPOCH}_{trainer.epoch + 1}' + else: + return f'{LogKeys.ITER}_{trainer.iter + 1}' + + def _do_save(self, trainer, save_strategy): + # prefix like 'epoch-1' or 'iter-1' + prefix = self.generate_prefix(trainer, save_strategy) + if self.processor.should_save_on_rank(trainer): + if is_master(): + if save_strategy == CheckpointStrategy.by_epoch: + self.logger.info( + f'Saving checkpoint at {trainer.epoch + 1} epoch') + else: + self.logger.info( + f'Saving checkpoint at {trainer.iter + 1} iter') + self._save_checkpoint(trainer, prefix) + if is_master() and self.push_to_hub: + if self.push_to_hub_future is not None and not self.push_to_hub_future.done( + ): + self.logger.error( + f'Another uploading is running, ' + f'this uploading with message {prefix} will be canceled.') + return + self.push_to_hub_future = self._push_to_hub(trainer, prefix) + + def after_train_epoch(self, trainer): + if self.save_strategy != CheckpointStrategy.by_epoch: + return + + if self._should_save(trainer): + self._do_save(trainer, CheckpointStrategy.by_epoch) + + def after_train_iter(self, trainer): + if self.save_strategy != CheckpointStrategy.by_step: + return + + if self._should_save(trainer): + self._do_save(trainer, CheckpointStrategy.by_step) + + def after_run(self, trainer): + if self.push_to_hub_future is not None and not self.push_to_hub_future.done( + ): + self.logger.info('Train finished. Uploading models, waiting...') + while not self.push_to_hub_future.done(): + time.sleep(1) + self.logger.info('Uploading models done.') + + def _push_to_hub(self, trainer, prefix): + if self.is_model_id is None: + self.is_model_id = check_model_is_id(trainer.input_model_id, + self.hub_token) + self.tag += 1 + return push_to_hub_async( + self.hub_repo_id, + self.output_dir, + token=self.hub_token, + private=self.private_hub, + commit_message=prefix, + tag=f'v1.{self.tag}', + revision=self.hub_revision, + source_repo=trainer.input_model_id if self.is_model_id else '') + + def save_evaluate_results(self, trainer): + with open(os.path.join(self.output_dir, self.EVAL_RESULT_FILE), + 'w') as f: + f.write(str(trainer.metric_values)) + + def _save_checkpoint(self, trainer, prefix): + """Save checkpoint files and remove obsolete ones + """ + checkpoint_path_prefix = os.path.join(self.save_dir, prefix) + meta = self._create_training_state(trainer) + self.processor.save_checkpoints(trainer, checkpoint_path_prefix, + self.output_dir, meta) + self.save_evaluate_results(trainer) + self.history_checkpoints.append(checkpoint_path_prefix) + self._remove_obsolete_checkpoints(trainer) + return prefix + + def _remove_obsolete_checkpoints(self, trainer): + if self.max_checkpoint_num is not None and \ + len(self.history_checkpoints) > self.max_checkpoint_num: + history_checkpoints = [ckpt for ckpt in self.history_checkpoints] + self.history_checkpoints.clear() + for i, checkpoint_path_prefix in enumerate(history_checkpoints): + if i < len(history_checkpoints) - self.max_checkpoint_num: + self.logger.info( + f'deleting checkpoint: {checkpoint_path_prefix}') + self.processor.remove_checkpoints( + trainer, checkpoint_path_prefix=checkpoint_path_prefix) + else: + self.history_checkpoints.append(checkpoint_path_prefix) + + def _should_save(self, trainer): + if self.save_strategy == CheckpointStrategy.by_epoch: + check_last = self.is_last_epoch + check_frequency = self.every_n_epochs + elif self.save_strategy == CheckpointStrategy.by_step: + check_last = self.is_last_iter + check_frequency = self.every_n_iters + else: + return False + + if check_frequency(trainer, + self.interval) or (self.save_last + and check_last(trainer)): + return True + return False + + def _create_training_state(self, trainer): + self.rng_state = { + 'random': random.getstate(), + 'numpy': np.random.get_state(), + 'cpu': torch.random.get_rng_state(), + 'cuda': torch.cuda.get_rng_state_all(), + } + + # keep epoch/iter/inner_iter/random_state + meta = { + 'epoch': trainer.epoch, + 'iter': trainer.iter + 1, + 'inner_iter': trainer.inner_iter + 1, + 'rng_state': self.rng_state, + } + + # keep hooks state + i = 0 + for hook in trainer.hooks: + if hasattr(hook, 'state_dict') and getattr(hook, '_should_save', + True): + meta[f'{hook.__class__}-{i}'] = hook.state_dict() + i += 1 + + return meta + + +@HOOKS.register_module(module_name=Hooks.BestCkptSaverHook) +class BestCkptSaverHook(CheckpointHook): + """ + Save best checkpoints hook. + + Args: + metric_key (str): Metric key to compare rule for best score. + save_best(bool): Save the best checkpoint, if set to False, this hook will have no effect. + rule (str): Comparison rule for best score. Support "max" and "min". If rule is "max", the checkpoint + at the maximum `metric_key` will be saved, If rule is "min", the checkpoint at the minimum `metric_key` + will be saved. + save_file_name: The manual specified saving file name. + restore_best (bool): Whether to restore the best checkpoint after training. + max_checkpoint_num (int): The max number of checkpoint files, default None which means never delete anything. + If the number exceeding the limit, checkpoints with worse metric will be deleted, which is judged by the + `rule` and `metric_key` arguments. + + The `BestCkptSaverHook` class accepts `output_sub_dir` and `output_dir` argument as its super class do. + If neither of them are passed, the default value is `{save_dir}/output_best`. + + This class will not accept the `interval` or `save_strategy` or `by_epoch` argument, because the saving interval + will follow the `EvaluationHook`. + """ + + PRIORITY = Priority.LOW + rule_map = {'max': lambda x, y: x > y, 'min': lambda x, y: x < y} + + def __init__(self, + metric_key: str, + save_best: Optional[bool] = True, + rule: Optional[str] = 'max', + save_file_name: Optional[str] = None, + restore_best: Optional[bool] = False, + max_checkpoint_num: Optional[int] = 1, + **kwargs): + assert rule in ['max', 'min'], 'Only support "max" or "min" rule now.' + output_kwargs = {} + if 'output_sub_dir' not in kwargs and 'output_dir' not in kwargs: + output_kwargs['output_sub_dir'] = ModelFile.TRAIN_BEST_OUTPUT_DIR + kwargs.pop('interval', None) + kwargs.pop('save_strategy', None) + super().__init__( + max_checkpoint_num=max_checkpoint_num, + **kwargs, + **output_kwargs, + ) + self.save_best = save_best + self.metric_key = metric_key + self.rule = rule + self._best_metric = None + self._best_ckpt_file = None + self.save_file_name = save_file_name + self.restore_best = restore_best + self.history_checkpoints = set() + + def after_train_epoch(self, trainer): + from modelscope.trainers.hooks import EvaluationHook + eval_hook = trainer.get_hook(EvaluationHook) + if len(eval_hook) == 0: + self.logger.error( + 'Trying to save the best checkpoint, but there is no evaluation, skipping.' + ) + + if eval_hook[0].last_eval_tag == ( + 'epoch', trainer.epoch) and self._should_save(trainer): + self._do_save(trainer, 'by_epoch') + + def after_train_iter(self, trainer): + from modelscope.trainers.hooks import EvaluationHook + eval_hook = trainer.get_hook(EvaluationHook) + if len(eval_hook) == 0: + self.logger.error( + 'Trying to save the best checkpoint, but there is no evaluation, skipping.' + ) + + if eval_hook[0].last_eval_tag == ( + 'iter', trainer.iter) and self._should_save(trainer): + self._do_save(trainer, 'by_step') + + def _should_save(self, trainer): + return self._is_best_metric(trainer.metric_values) and self.save_best + + def _is_best_metric(self, metric_values): + if metric_values is None: + return False + + if self.metric_key not in metric_values: + raise ValueError( + f'Not find metric_key: {self.metric_key} in {metric_values}') + + if self._best_metric is None: + self._best_metric = metric_values[self.metric_key] + return True + else: + compare_fn = self.rule_map[self.rule] + if compare_fn(metric_values[self.metric_key], self._best_metric): + self._best_metric = metric_values[self.metric_key] + return True + return False + + def generate_prefix(self, trainer, save_strategy): + if save_strategy == CheckpointStrategy.by_epoch: + return f'best_{LogKeys.EPOCH}{trainer.epoch + 1}_{self.metric_key}{self._best_metric}' + else: + return f'best_{LogKeys.ITER}{trainer.iter + 1}_{self.metric_key}{self._best_metric}' + + def _save_checkpoint(self, trainer, prefix): + checkpoint_path_prefix = self.save_file_name + if checkpoint_path_prefix is None: + checkpoint_path_prefix = os.path.join(self.save_dir, prefix) + else: + checkpoint_path_prefix = os.path.join(self.save_dir, + checkpoint_path_prefix) + + self._best_ckpt_file = checkpoint_path_prefix + meta = self._create_training_state(trainer) + self.processor.save_checkpoints(trainer, checkpoint_path_prefix, + self.output_dir, meta) + self.save_evaluate_results(trainer) + self.history_checkpoints.add(checkpoint_path_prefix) + self._remove_obsolete_checkpoints(trainer) + return prefix + + def _remove_obsolete_checkpoints(self, trainer): + + def extract_metric_from_filename(name1): + metric1 = float(name1.split(self.metric_key)[1]) + if self.rule == 'max': + return -metric1 + else: + return metric1 + + if self.max_checkpoint_num is not None and \ + len(self.history_checkpoints) > self.max_checkpoint_num: + history_checkpoints = sorted( + self.history_checkpoints, key=extract_metric_from_filename) + self.history_checkpoints.clear() + for i, checkpoint_path_prefix in enumerate(history_checkpoints): + if i < self.max_checkpoint_num: + self.history_checkpoints.add(checkpoint_path_prefix) + else: + self.logger.info( + f'deleting checkpoint: {checkpoint_path_prefix}') + self.processor.remove_checkpoints( + trainer, checkpoint_path_prefix=checkpoint_path_prefix) + + def state_dict(self): + return { + 'best_metric': self._best_metric, + } + + def load_state_dict(self, state_dict): + if state_dict is not None and len(state_dict) > 0: + self._best_metric = state_dict.get('best_metric') + else: + self.logger.warning( + 'The state_dict is not available, the best metric value will be affected.' + ) + + def after_run(self, trainer): + if self.restore_best: + # If restore_best is True, will call the LoadCheckpointHook to load the best checkpoint + # for later evaluation or prediction. + from modelscope.trainers.hooks.checkpoint.load_checkpoint_hook import LoadCheckpointHook + LoadCheckpointHook.load_checkpoint(self._best_ckpt_file, trainer) diff --git a/modelscope/trainers/hooks/checkpoint/checkpoint_processor.py b/modelscope/trainers/hooks/checkpoint/checkpoint_processor.py new file mode 100644 index 00000000..f28fc397 --- /dev/null +++ b/modelscope/trainers/hooks/checkpoint/checkpoint_processor.py @@ -0,0 +1,276 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +import re +import shutil + +from modelscope.metainfo import Pipelines +from modelscope.utils.checkpoint import (load_checkpoint, save_checkpoint, + save_configuration) +from modelscope.utils.constant import ModelFile +from modelscope.utils.logger import get_logger +from modelscope.utils.torch_utils import is_master + + +class CheckpointProcessor: + + TRAINER_STATE_SUFFIX = '_trainer_state.pth' + + MODEL_STATE_SUFFIX = '.pth' + + def prepare_output(self, trainer, output_dir): + """Prepares the output of target folder. + + This is a strategic function which can be registered by other hook's function. + + Args: + trainer: The trainer instance. + output_dir: The target folder used in inference. + """ + model = trainer.unwrap_module(trainer.model) + config = trainer.cfg + + # override pipeline by tasks name after finetune done, + # avoid case like fill mask pipeline with a text cls task + if config['task'] in [ + getattr(Pipelines, attr) for attr in dir(Pipelines) + if not attr.startswith('__') + ]: + # TODO a temp fix to avoid pipeline_name and task mismatch + config['pipeline'] = {'type': config['task']} + + self.copy_files_and_dump_config(trainer, output_dir, config, + self._bin_file(model)) + + @staticmethod + def copy_files_and_dump_config(trainer, output_dir, config, bin_file): + """Copy useful files to target output folder and dumps the target configuration.json. + """ + model = trainer.unwrap_module(trainer.model) + + class SaveConfig: + + def __init__(self, output_dir, config): + self.output_dir = output_dir + self.config = config + + def __call__(self, _output_dir, _config): + self.config = _config + + def save_config(self): + save_configuration(self.output_dir, self.config) + + for pop_key in [ + 'push_to_hub', 'hub_repo_id', 'hub_token', 'private_hub' + ]: + if config.safe_get('train.checkpoint.period.' + + pop_key) is not None: + config.safe_get('train.checkpoint.period').pop(pop_key) + if config.safe_get('train.checkpoint.best.' + pop_key) is not None: + config.safe_get('train.checkpoint.best').pop(pop_key) + + save_config_fn = SaveConfig(output_dir, config) + + if hasattr(model, 'save_pretrained'): + # Save pretrained of model, skip saving checkpoint + model.save_pretrained( + output_dir, + bin_file, + save_function=lambda *args, **kwargs: None, + config=save_config_fn.config, + save_config_function=save_config_fn) + + if trainer.train_preprocessor is not None: + trainer.train_preprocessor.save_pretrained( + output_dir, + save_config_fn.config, + save_config_function=save_config_fn) + if trainer.eval_preprocessor is not None: + trainer.eval_preprocessor.save_pretrained( + output_dir, + save_config_fn.config, + save_config_function=save_config_fn) + save_config_fn.save_config() + + @staticmethod + def _bin_file(model): + """Get bin file path. + """ + default_bin_file = ModelFile.TORCH_MODEL_BIN_FILE + if hasattr(model, + 'model_dir') and ModelFile.TORCH_MODEL_FILE in os.listdir( + model.model_dir): + default_bin_file = ModelFile.TORCH_MODEL_FILE + return default_bin_file + + def save_checkpoints(self, + trainer, + checkpoint_path_prefix, + output_dir, + meta=None): + """Save the state dict for trainer and model. + + This is a strategic function which can be registered by other hook's function. + + Args: + trainer(`EpochBasedTrainer`): The trainer instance. + checkpoint_path_prefix(`str`): The saving dir with a prefix. + like: /tmp/test/epoch_0 + output_dir(`str`): The output dir for inference. + meta: (`dict`): The meta info needed to be saved into files. + """ + model = trainer.unwrap_module(trainer.model) + _model_file, _train_state_file = self._get_state_file_name( + checkpoint_path_prefix) + + # Save pth file without model state_dict + self.save_trainer_state(trainer, model, _train_state_file, meta) + self.save_model_state(model, _model_file) + self.link(model, _model_file, output_dir) + + def remove_checkpoints(self, trainer, checkpoint_path_prefix): + """Remove obsolete checkpoint files. + + This is a strategic function which can be registered by other hook's function. + + Args: + trainer(`EpochBasedTrainer`): The trainer instance. + checkpoint_path_prefix(`str`): The saving dir with a prefix. + like: /tmp/test/epoch_0 + """ + _model_file, _train_state_file = self._get_state_file_name( + checkpoint_path_prefix) + if os.path.isfile(_train_state_file): + os.remove(_train_state_file) + + if os.path.isfile(_model_file): + os.remove(_model_file) + + def should_save_on_rank(self, trainer): + """Used in ddp or other distributed training scenario, returns whether do saving in current rank. + + This is a strategic function which can be registered by other hook's function. + + Args: + trainer(`EpochBasedTrainer`): The trainer instance. + """ + return is_master() + + def link(self, model, src_file, output_dir): + """Links the src bin file to the output folder. + + Args: + model: The model instance. + src_file: The src bin file path. + output_dir: The target folder used in inference. + """ + + bin_file = self._bin_file(model) + dest_file = os.path.join(output_dir, bin_file) + if os.path.isfile(dest_file): + os.unlink(dest_file) + + try: + os.link(src_file, dest_file) + except OSError as e: + get_logger().error( + f'Link {src_file} to {dest_file} error: {e}, ' + 'changing to copy the bin file, this may use more disk space.') + shutil.copyfile(src_file, dest_file) + + def save_trainer_state(self, trainer, model, train_state_file, meta): + """Save the trainer state, including optimizer/lr_scheduler's state dict, random states etc. + + Args: + trainer: The trainer instance. + model: The model instance. + train_state_file: The target file name for saving trainer states. + meta: Some extra meta info. + """ + save_checkpoint( + model, + train_state_file, + trainer.optimizer, + trainer.lr_scheduler, + meta=meta, + with_model=False) + + def save_model_state(self, model, model_file): + """Save the model state. + + Args: + model: The model instance. + model_file: The target file name for saving model states. + """ + save_checkpoint( + model, model_file, None, None, meta=None, with_meta=False) + + def load_checkpoints(self, checkpoint_path_prefix, trainer, load_all_state, + strict): + """Load checkpoint files of trainer state and model state. + + This is a strategic function which can be registered by other hook's function. + + Args: + checkpoint_path_prefix(str): The checkpoint dir with prefix or a model state file. + Example: '/tmp/test/epoch_0' or '/tmp/test/epoch_0.pth' + trainer(`EpochBasedTrainer`): The trainer instance. + load_all_state(`boolean`): Load all states (else load only module states). + strict(`boolean`): If strict, any unmatched keys will cause an error. + + Returns: + The meta info in json. + """ + _model_file, _train_state_file = self._get_state_file_name( + checkpoint_path_prefix) + meta = {} + if os.path.isfile(_train_state_file): + meta = self.load_trainer_state(trainer, _train_state_file, + load_all_state) + else: + print(f'No trainer state file {_train_state_file} found, skip.') + self.load_model_state(trainer, _model_file, strict) + return meta + + @staticmethod + def load_trainer_state(trainer, train_state_file, load_all_state): + """Load trainer state file. + """ + + optimizer = getattr(trainer, 'optimizer', + None) if load_all_state else None + lr_scheduler = getattr(trainer, 'lr_scheduler', + None) if load_all_state else None + return load_checkpoint(train_state_file, None, optimizer, lr_scheduler) + + def load_model_state(self, trainer, model_file, strict): + """Load model state file. + """ + return load_checkpoint(model_file, + trainer.unwrap_module(trainer.model), None, + None) + + @staticmethod + def _get_state_file_name(checkpoint_path_prefix): + """Get the default file name for state files. + + If the input is a checkpoint dir with prefix, this function will append suffix for both checkpoint files. + If the input is an absolute file name, this function will return it as the model file name, and append + suffix for the trainer file name. + + NOTE: a best checkpoint filename with float or int metric value inside + will not be judged as having a extension file name. like: '/tmp/test/epoch_0_accuracy0.85' + + Args: + checkpoint_path_prefix(`str`): The checkpoint dir with prefix or a model state file + with extension file name. like: '/tmp/test/epoch_0' + + Returns: + A tuple of model state file name and trainer state file name. + """ + base, ext = os.path.splitext(checkpoint_path_prefix) + if len(ext) == 0 or re.match(r'^\d+$', ext[1:]): + return checkpoint_path_prefix + CheckpointProcessor.MODEL_STATE_SUFFIX, \ + checkpoint_path_prefix + CheckpointProcessor.TRAINER_STATE_SUFFIX # noqa + else: + return checkpoint_path_prefix, base + CheckpointProcessor.TRAINER_STATE_SUFFIX.split( + '.')[0] + '.' + ext[1:] diff --git a/modelscope/trainers/hooks/checkpoint/load_checkpoint_hook.py b/modelscope/trainers/hooks/checkpoint/load_checkpoint_hook.py new file mode 100644 index 00000000..3ccb800f --- /dev/null +++ b/modelscope/trainers/hooks/checkpoint/load_checkpoint_hook.py @@ -0,0 +1,138 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import random +from typing import Optional + +import numpy as np +import torch +from packaging import version + +from modelscope.metainfo import Hooks +from modelscope.trainers.hooks.builder import HOOKS +from modelscope.trainers.hooks.checkpoint.checkpoint_processor import \ + CheckpointProcessor +from modelscope.trainers.hooks.hook import Hook +from modelscope.trainers.hooks.priority import Priority +from modelscope.utils.logger import get_logger + + +@HOOKS.register_module(module_name=Hooks.LoadCheckpointHook) +class LoadCheckpointHook(Hook): + """Load a checkpoint file at the beginning of training or evaluating. + + This hook does not need to be configured or saved in the config file. + User should use it by: + >>> trainer.train('some-checkpoint', load_all_state=True) + or + >>> trainer.evaluate('some-checkpoint') + instead. + + Args: + checkpoint_file (str): The checkpoint file to be loaded. + load_all_state (bool): Load all states(optimizer, epoch, lr_scheduler, random_state, etc.) when loading old + training state file or not. The model's state dict will only be loaded if False. + strict (bool): If strict, any unmatched keys will cause an error. + """ + + PRIORITY = Priority.HIGH + + _should_save = False + + # From 1.3.1 version we split one pth file to two files: trainer state pth file/model state pth file. + _TWO_PTH_FILE_VERSION = '1.3.1' + + def __init__( + self, + checkpoint_file: Optional[str] = None, + load_all_state: Optional[bool] = True, + strict: Optional[bool] = False, + ): + self.checkpoint_file = checkpoint_file + self.rng_state = None + self.need_load_rng_state = False + self.load_all_state = load_all_state + self.strict = strict + self.processor = CheckpointProcessor() + + def before_run(self, trainer): + if not hasattr(trainer, 'logger'): + self.logger = get_logger() + else: + self.logger = trainer.logger + + if self.checkpoint_file is not None: + meta = self.load_checkpoint(self.checkpoint_file, trainer, + self.load_all_state, self.strict) + self.rng_state = meta.get('rng_state') + self.need_load_rng_state = self.load_all_state + + def before_train_iter(self, trainer): + if self.need_load_rng_state: + if self.rng_state is not None: + random.setstate(self.rng_state['random']) + np.random.set_state(self.rng_state['numpy']) + torch.random.set_rng_state(self.rng_state['cpu']) + if torch.cuda.is_available(): + torch.cuda.random.set_rng_state_all(self.rng_state['cuda']) + self.need_load_rng_state = False + else: + self.logger.info( + 'Random state cannot be found in checkpoint file, ' + 'this may cause a random data order or model initialization.' + ) + + @staticmethod + def _restore_training_state(trainer, meta): + trainer._epoch = meta.get('epoch', trainer._epoch) + trainer._iter = meta.get('iter', trainer._iter) + trainer._inner_iter = meta.get('inner_iter', trainer._inner_iter) + + i = 0 + for hook in trainer.hooks: + if hasattr(hook, 'load_state_dict') and getattr( + hook, '_should_save', True): + key = f'{hook.__class__}-{i}' + if key in meta: + hook.load_state_dict(meta.get(key, {})) + else: + trainer.logger.warning( + f'The state_dict of hook {hook.__class__} at index {i} is not found in the checkpoint file.' + ) + i += 1 + + @classmethod + def load_checkpoint(cls, + filename, + trainer, + load_all_state=True, + strict=False): + """A static method to load checkpoint files. + + Args: + filename(str): An absolute model bin file(pth or bin) or a dir path with a file prefix(like epoch_1). + trainer(`EpochBasedTrainer`): The trainer instance. + load_all_state(`bool`): Load all states including the trainer states. + strict(`bool`): Load module state dict strictly. + + Returns: + A dict containing the train states saved by `_create_training_state` + """ + meta = cls().processor.load_checkpoints(filename, trainer, + load_all_state, strict) + if load_all_state: + cls._restore_training_state(trainer, meta) + + if meta is not None: + _version = meta.get('modelscope') + if _version is not None and version.parse( + _version) < version.parse( + LoadCheckpointHook._TWO_PTH_FILE_VERSION): + trainer.logger.warning( + 'The unique pth file is split into a model file and ' + f'a trainer file since version {LoadCheckpointHook._TWO_PTH_FILE_VERSION},' + 'consider re-training your model or ' + 'using a converting script to split the single pth file into two.' + ) + trainer.logger.info( + f'Checkpoint {filename} saving time: {meta.get("time")}, modelscope version: {_version}' + ) + return meta diff --git a/modelscope/trainers/hooks/checkpoint_hook.py b/modelscope/trainers/hooks/checkpoint_hook.py deleted file mode 100644 index 59832105..00000000 --- a/modelscope/trainers/hooks/checkpoint_hook.py +++ /dev/null @@ -1,749 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. -import os -import random -import re -import time - -import numpy as np -import torch -from packaging import version - -from modelscope.hub.check_model import check_model_is_id -from modelscope.hub.push_to_hub import push_to_hub_async -from modelscope.metainfo import Hooks, Pipelines -from modelscope.utils.checkpoint import (load_checkpoint, save_checkpoint, - save_configuration) -from modelscope.utils.constant import LogKeys, ModelFile -from modelscope.utils.logger import get_logger -from modelscope.utils.torch_utils import is_master -from .builder import HOOKS -from .hook import Hook -from .priority import Priority - - -@HOOKS.register_module(module_name=Hooks.CheckpointHook) -class CheckpointHook(Hook): - """Save checkpoints periodically. - - Args: - interval (int): The frequency to save model. If `by_epoch=True`, - it means the number of epochs, else means the number of iterations - by_epoch (bool): Saving checkpoints by epoch or by iteration. - save_optimizer (bool): Whether to save optimizer state dict. Default: True. - save_dir (str): The directory to save checkpoints. If is None, use `trainer.work_dir` - output_sub_dir (str): The sub folder under the `save_dir` to save the output checkpoint for inference. - Default 'output'. - save_last (bool): Whether to save the last checkpoint. Default: True. - max_checkpoint_num (int): The max number of checkpoint files, default None which means never delete anything. - If the number exceeding the limit, earlier checkpoints will be deleted first. - """ - - PRIORITY = Priority.LOW - - TRAINER_STATE_SUFFIX = '_trainer_state.pth' - - MODEL_STATE_SUFFIX = '.pth' - - def __init__(self, - interval=0, - by_epoch=True, - save_optimizer=True, - save_dir=None, - output_sub_dir=ModelFile.TRAIN_OUTPUT_DIR, - save_last=True, - max_checkpoint_num=None, - push_to_hub=False, - model_id_with_org=None, - hub_token=None, - private_hub=True, - **kwargs): - self.interval = interval - self.by_epoch = by_epoch - self.save_optimizer = save_optimizer - self.save_dir = save_dir - self.output_sub_dir = output_sub_dir - self.save_last = save_last - self.rng_state = None - self.max_checkpoint_num = None - self.push_to_hub = push_to_hub - self.model_id_with_org = model_id_with_org - self.hub_token = hub_token - self.private_hub = private_hub - self.is_model_id = None - self.push_to_hub_future = None - if max_checkpoint_num is not None: - self.max_checkpoint_num = max(int(max_checkpoint_num), 1) - self.history_checkpoints = [] - - def before_run(self, trainer): - if not self.save_dir: - self.save_dir = trainer.work_dir - - if not os.path.exists(self.save_dir): - os.makedirs(self.save_dir, exist_ok=True) - - if not hasattr(trainer, 'logger'): - self.logger = get_logger() - else: - self.logger = trainer.logger - - if is_master(): - output_dir = os.path.join(self.save_dir, self.output_sub_dir) - # only global master prepares the output folder - self.prepare_output(trainer, output_dir) - self.logger.info(f'Checkpoints will be saved to {self.save_dir}') - - def generate_prefix(self, trainer): - if self.by_epoch: - return f'{LogKeys.EPOCH}_{trainer.epoch + 1}' - else: - return f'{LogKeys.ITER}_{trainer.iter + 1}' - - def after_train_epoch(self, trainer): - if not self.by_epoch: - return - - if self._should_save(trainer): - # prefix like 'epoch-1' or 'iter-1' - prefix = self.generate_prefix(trainer) - if self.should_save_on_rank(trainer): - if is_master(): - self.logger.info( - f'Saving checkpoint at {trainer.epoch + 1} epoch') - self._save_checkpoint(trainer, prefix) - if is_master() and self.push_to_hub: - if self.push_to_hub_future is not None and not self.push_to_hub_future.done( - ): - self.logger.error( - f'Another uploading is running, ' - f'this uploading with message {prefix} will be canceled.' - ) - return - self.push_to_hub_future = self._push_to_hub(trainer, prefix) - - def after_train_iter(self, trainer): - if self.by_epoch: - return - - if self._should_save(trainer): - # prefix like 'epoch-1' or 'iter-1' - prefix = self.generate_prefix(trainer) - if self.should_save_on_rank(trainer): - if is_master(): - self.logger.info( - f'Saving checkpoint at {trainer.iter + 1} iter') - self._save_checkpoint(trainer, prefix) - if is_master() and self.push_to_hub: - if self.push_to_hub_future is not None and not self.push_to_hub_future.done( - ): - self.logger.error( - f'Another uploading is running, ' - f'this uploading with message {prefix} will be canceled.' - ) - return - self.push_to_hub_future = self._push_to_hub(trainer, prefix) - - def after_run(self, trainer): - if self.push_to_hub_future is not None and not self.push_to_hub_future.done( - ): - self.logger.info('Train finished. Uploading models, waiting...') - while not self.push_to_hub_future.done(): - time.sleep(1) - self.logger.info('Uploading models done.') - - def _push_to_hub(self, trainer, prefix): - if self.is_model_id is None: - self.is_model_id = check_model_is_id(trainer.input_model_id, - self.hub_token) - - return push_to_hub_async( - self.model_id_with_org, - os.path.join(self.save_dir, self.output_sub_dir), - token=self.hub_token, - private=self.private_hub, - commit_message=prefix, - source_repo=trainer.input_model_id if self.is_model_id else '') - - def _save_checkpoint(self, trainer, prefix): - """Save checkpoint files and remove obsolete ones - """ - checkpoint_path_prefix = os.path.join(self.save_dir, prefix) - meta = self._create_training_state(trainer) - self.save_checkpoints(trainer, checkpoint_path_prefix, - self.output_sub_dir, meta) - self.history_checkpoints.append(checkpoint_path_prefix) - self._remove_obsolete_checkpoints(trainer) - return prefix - - def _remove_obsolete_checkpoints(self, trainer): - if self.max_checkpoint_num is not None and \ - len(self.history_checkpoints) > self.max_checkpoint_num: - history_checkpoints = [ckpt for ckpt in self.history_checkpoints] - self.history_checkpoints.clear() - for i, checkpoint_path_prefix in enumerate(history_checkpoints): - if i < len(history_checkpoints) - self.max_checkpoint_num: - self.logger.info( - f'deleting checkpoint: {checkpoint_path_prefix}') - self.remove_checkpoints( - trainer, checkpoint_path_prefix=checkpoint_path_prefix) - else: - self.history_checkpoints.append(checkpoint_path_prefix) - - def _should_save(self, trainer): - if self.by_epoch: - check_last = self.is_last_epoch - check_frequency = self.every_n_epochs - else: - check_last = self.is_last_iter - check_frequency = self.every_n_iters - - if check_frequency(trainer, - self.interval) or (self.save_last - and check_last(trainer)): - return True - return False - - def _create_training_state(self, trainer): - self.rng_state = { - 'random': random.getstate(), - 'numpy': np.random.get_state(), - 'cpu': torch.random.get_rng_state(), - 'cuda': torch.cuda.get_rng_state_all(), - } - - # keep epoch/iter/inner_iter/random_state - meta = { - 'epoch': trainer.epoch, - 'iter': trainer.iter + 1, - 'inner_iter': trainer.inner_iter + 1, - 'rng_state': self.rng_state, - } - - # keep hooks state - i = 0 - for hook in trainer.hooks: - if hasattr(hook, 'state_dict') and getattr(hook, '_should_save', - True): - meta[f'{hook.__class__}-{i}'] = hook.state_dict() - i += 1 - - return meta - - @staticmethod - def copy_files_and_dump_config(trainer, output_dir, config, bin_file): - """Copy useful files to target output folder and dumps the target configuration.json. - """ - model = trainer.unwrap_module(trainer.model) - - class SaveConfig: - - def __init__(self, output_dir, config): - self.output_dir = output_dir - self.config = config - - def __call__(self, _output_dir, _config): - self.config = _config - - def save_config(self): - save_configuration(self.output_dir, self.config) - - for pop_key in [ - 'push_to_hub', 'model_id_with_org', 'hub_token', 'private_hub' - ]: - if config.safe_get('train.checkpoint.period.' - + pop_key) is not None: - config.safe_get('train.checkpoint.period').pop(pop_key) - if config.safe_get('train.checkpoint.best.' + pop_key) is not None: - config.safe_get('train.checkpoint.best').pop(pop_key) - - save_config_fn = SaveConfig(output_dir, config) - - if hasattr(model, 'save_pretrained'): - # Save pretrained of model, skip saving checkpoint - model.save_pretrained( - output_dir, - bin_file, - save_function=lambda *args, **kwargs: None, - config=save_config_fn.config, - save_config_function=save_config_fn) - - if trainer.train_preprocessor is not None: - trainer.train_preprocessor.save_pretrained( - output_dir, - save_config_fn.config, - save_config_function=save_config_fn) - if trainer.eval_preprocessor is not None: - trainer.eval_preprocessor.save_pretrained( - output_dir, - save_config_fn.config, - save_config_function=save_config_fn) - save_config_fn.save_config() - - @staticmethod - def _bin_file(model): - """Get bin file path. - """ - default_bin_file = ModelFile.TORCH_MODEL_BIN_FILE - if hasattr(model, - 'model_dir') and ModelFile.TORCH_MODEL_FILE in os.listdir( - model.model_dir): - default_bin_file = ModelFile.TORCH_MODEL_FILE - return default_bin_file - - @Hook.overload_func(name='CheckpointHook.prepare_output') - def prepare_output(self, trainer, output_dir): - """Prepares the output of target folder. - - This is a strategic function which can be registered by other hook's function. - - Args: - trainer: The trainer instance. - output_dir: The target folder used in inference. - """ - model = trainer.unwrap_module(trainer.model) - config = trainer.cfg - - # override pipeline by tasks name after finetune done, - # avoid case like fill mask pipeline with a text cls task - if config['task'] in [ - getattr(Pipelines, attr) for attr in dir(Pipelines) - if not attr.startswith('__') - ]: - # TODO a temp fix to avoid pipeline_name and task mismatch - config['pipeline'] = {'type': config['task']} - - self.copy_files_and_dump_config(trainer, output_dir, config, - self._bin_file(model)) - - def link(self, model, src_file, output_dir): - """Links the src bin file to the output folder. - - Args: - model: The model instance. - src_file: The src bin file path. - output_dir: The target folder used in inference. - """ - - bin_file = self._bin_file(model) - dest_file = os.path.join(output_dir, bin_file) - if os.path.isfile(dest_file): - os.unlink(dest_file) - - os.link(src_file, dest_file) - - def save_trainer_state(self, trainer, model, train_state_file, meta): - """Save the trainer state, including optimizer/lr_scheduler's state dict, random states etc. - - Args: - trainer: The trainer instance. - model: The model instance. - train_state_file: The target file name for saving trainer states. - meta: Some extra meta info. - """ - save_checkpoint( - model, - train_state_file, - trainer.optimizer, - trainer.lr_scheduler, - meta=meta, - with_model=False) - - def save_model_state(self, model, model_file): - """Save the model state. - - Args: - model: The model instance. - model_file: The target file name for saving model states. - """ - save_checkpoint( - model, model_file, None, None, meta=None, with_meta=False) - - @Hook.overload_func(name='CheckpointHook.save_checkpoints') - def save_checkpoints(self, - trainer, - checkpoint_path_prefix, - output_sub_dir, - meta=None): - """Save the state dict for trainer and model. - - This is a strategic function which can be registered by other hook's function. - - Args: - trainer(`EpochBasedTrainer`): The trainer instance. - checkpoint_path_prefix(`str`): The saving dir with a prefix. - like: /tmp/test/epoch_0 - output_sub_dir(`str`): The sub-dir in the saving dir used in inference. - meta: (`dict`): The meta info needed to be saved into files. - """ - model = trainer.unwrap_module(trainer.model) - _model_file, _train_state_file = _get_state_file_name( - checkpoint_path_prefix) - - # Save pth file without model state_dict - self.save_trainer_state(trainer, model, _train_state_file, meta) - self.save_model_state(model, _model_file) - output_dir = os.path.join(self.save_dir, output_sub_dir) - self.link(model, _model_file, output_dir) - - @Hook.overload_func(name='CheckpointHook.remove_checkpoints') - def remove_checkpoints(self, trainer, checkpoint_path_prefix): - """Remove obsolete checkpoint files. - - This is a strategic function which can be registered by other hook's function. - - Args: - trainer(`EpochBasedTrainer`): The trainer instance. - checkpoint_path_prefix(`str`): The saving dir with a prefix. - like: /tmp/test/epoch_0 - """ - _model_file, _train_state_file = _get_state_file_name( - checkpoint_path_prefix) - if os.path.isfile(_train_state_file): - os.remove(_train_state_file) - - if os.path.isfile(_model_file): - os.remove(_model_file) - - @Hook.overload_func(name='CheckpointHook.should_save_on_rank') - def should_save_on_rank(self, trainer): - """Used in ddp or other distributed training scenario, returns whether do saving in current rank. - - This is a strategic function which can be registered by other hook's function. - - Args: - trainer(`EpochBasedTrainer`): The trainer instance. - """ - return is_master() - - -@HOOKS.register_module(module_name=Hooks.BestCkptSaverHook) -class BestCkptSaverHook(CheckpointHook): - """ - Save best checkpoints hook. - - Args: - metric_key (str): Metric key to compare rule for best score. - rule (str): Comparison rule for best score. Support "max" and "min". If rule is "max", the checkpoint - at the maximum `metric_key` will be saved, If rule is "min", the checkpoint at the minimum `metric_key` - will be saved. - by_epoch (bool): Save best checkpoints by epoch or by iteration. - save_optimizer (bool): Whether to save optimizer state dict. Default: True. - save_dir (str): Output directory to save best checkpoint. - output_sub_dir (str): The sub folder under the `save_dir` to save the output checkpoint for inference. - Default 'output_best'. - restore_best (bool): Whether to restore the best checkpoint after training. - max_checkpoint_num (int): The max number of checkpoint files, default None which means never delete anything. - If the number exceeding the limit, checkpoints with worse metric will be deleted, which is judged by the - `rule` and `metric_key` arguments. - """ - - PRIORITY = Priority.LOW - rule_map = {'max': lambda x, y: x > y, 'min': lambda x, y: x < y} - - def __init__(self, - metric_key, - rule='max', - by_epoch=True, - save_optimizer=True, - save_dir=None, - output_sub_dir=ModelFile.TRAIN_BEST_OUTPUT_DIR, - save_file_name=None, - restore_best=False, - max_checkpoint_num=1, - interval=0, - **kwargs): - assert rule in ['max', 'min'], 'Only support "max" or "min" rule now.' - super().__init__( - interval=interval, - by_epoch=by_epoch, - save_optimizer=save_optimizer, - save_dir=save_dir, - output_sub_dir=output_sub_dir, - max_checkpoint_num=max_checkpoint_num, - **kwargs, - ) - self.metric_key = metric_key - self.rule = rule - self._best_metric = None - self._best_ckpt_file = None - self.save_file_name = save_file_name - self.restore_best = restore_best - self.history_checkpoints = set() - - def _should_save(self, trainer): - return self._is_best_metric(trainer.metric_values) - - def _is_best_metric(self, metric_values): - if metric_values is None: - return False - - if self.metric_key not in metric_values: - raise ValueError( - f'Not find metric_key: {self.metric_key} in {metric_values}') - - if self._best_metric is None: - self._best_metric = metric_values[self.metric_key] - return True - else: - compare_fn = self.rule_map[self.rule] - if compare_fn(metric_values[self.metric_key], self._best_metric): - self._best_metric = metric_values[self.metric_key] - return True - return False - - def generate_prefix(self, trainer): - if self.by_epoch: - return f'best_{LogKeys.EPOCH}{trainer.epoch + 1}_{self.metric_key}{self._best_metric}' - else: - return f'best_{LogKeys.ITER}{trainer.iter + 1}_{self.metric_key}{self._best_metric}' - - def _save_checkpoint(self, trainer, prefix): - checkpoint_path_prefix = self.save_file_name - if checkpoint_path_prefix is None: - checkpoint_path_prefix = os.path.join(self.save_dir, prefix) - else: - checkpoint_path_prefix = os.path.join(self.save_dir, - checkpoint_path_prefix) - - self._best_ckpt_file = checkpoint_path_prefix - meta = self._create_training_state(trainer) - self.save_checkpoints(trainer, checkpoint_path_prefix, - self.output_sub_dir, meta) - self.history_checkpoints.add(checkpoint_path_prefix) - self._remove_obsolete_checkpoints(trainer) - return prefix - - def _remove_obsolete_checkpoints(self, trainer): - - def extract_metric_from_filename(name1): - metric1 = float(name1.split(self.metric_key)[1]) - if self.rule == 'max': - return -metric1 - else: - return metric1 - - if self.max_checkpoint_num is not None and \ - len(self.history_checkpoints) > self.max_checkpoint_num: - history_checkpoints = sorted( - self.history_checkpoints, key=extract_metric_from_filename) - self.history_checkpoints.clear() - for i, checkpoint_path_prefix in enumerate(history_checkpoints): - if i < self.max_checkpoint_num: - self.history_checkpoints.add(checkpoint_path_prefix) - else: - self.logger.info( - f'deleting checkpoint: {checkpoint_path_prefix}') - self.remove_checkpoints( - trainer, checkpoint_path_prefix=checkpoint_path_prefix) - - def state_dict(self): - return { - 'best_metric': self._best_metric, - } - - def load_state_dict(self, state_dict): - if state_dict is not None and len(state_dict) > 0: - self._best_metric = state_dict.get('best_metric') - else: - self.logger.warning( - 'The state_dict is not available, the best metric value will be affected.' - ) - - def after_run(self, trainer): - if self.restore_best: - # If restore_best is True, will call the LoadCheckpointHook to load the best checkpoint - # for later evaluation or prediction. - LoadCheckpointHook.load_checkpoint(self._best_ckpt_file, trainer) - - -@HOOKS.register_module(module_name=Hooks.LoadCheckpointHook) -class LoadCheckpointHook(Hook): - """Load a checkpoint file at the beginning of training or evaluating. - - This hook does not need to be configured or saved in the config file. - User should use it by: - >>> trainer.train('some-checkpoint', load_all_state=True) - or - >>> trainer.evaluate('some-checkpoint') - instead. - - Args: - checkpoint_file (str): The checkpoint file to be loaded. - load_all_state (bool): Load all states(optimizer, epoch, lr_scheduler, random_state, etc.) when loading old - training state file or not. The model's state dict will only be loaded if False. - strict (bool): If strict, any unmatched keys will cause an error. - """ - - PRIORITY = Priority.HIGH - - _should_save = False - - _TWO_PTH_FILE_VERSION = '1.3.1' - - def __init__( - self, - checkpoint_file=None, - load_all_state=True, - strict=False, - ): - self.checkpoint_file = checkpoint_file - self.rng_state = None - self.need_load_rng_state = False - self.load_all_state = load_all_state - self.strict = strict - - def before_run(self, trainer): - if not hasattr(trainer, 'logger'): - self.logger = get_logger() - else: - self.logger = trainer.logger - - if self.checkpoint_file is not None: - meta = self.load_checkpoint(self.checkpoint_file, trainer, - self.load_all_state, self.strict) - self.rng_state = meta.get('rng_state') - self.need_load_rng_state = self.load_all_state - - def before_train_iter(self, trainer): - if self.need_load_rng_state: - if self.rng_state is not None: - random.setstate(self.rng_state['random']) - np.random.set_state(self.rng_state['numpy']) - torch.random.set_rng_state(self.rng_state['cpu']) - if torch.cuda.is_available(): - torch.cuda.random.set_rng_state_all(self.rng_state['cuda']) - self.need_load_rng_state = False - else: - self.logger.info( - 'Random state cannot be found in checkpoint file, ' - 'this may cause a random data order or model initialization.' - ) - - @staticmethod - def _restore_training_state(trainer, meta): - trainer._epoch = meta.get('epoch', trainer._epoch) - trainer._iter = meta.get('iter', trainer._iter) - trainer._inner_iter = meta.get('inner_iter', trainer._inner_iter) - - i = 0 - for hook in trainer.hooks: - if hasattr(hook, 'load_state_dict') and getattr( - hook, '_should_save', True): - key = f'{hook.__class__}-{i}' - if key in meta: - hook.load_state_dict(meta.get(key, {})) - else: - trainer.logger.warning( - f'The state_dict of hook {hook.__class__} at index {i} is not found in the checkpoint file.' - ) - i += 1 - - @classmethod - def load_checkpoint(cls, - filename, - trainer, - load_all_state=True, - strict=False): - """A static method to load checkpoint files. - - Args: - filename(str): An absolute model bin file(pth or bin) or a dir path with a file prefix(like epoch_1). - trainer(`EpochBasedTrainer`): The trainer instance. - load_all_state(`bool`): Load all states including the trainer states. - strict(`bool`): Load module state dict strictly. - - Returns: - A dict containing the train states saved by `_create_training_state` - """ - meta = cls().load_checkpoints(filename, trainer, load_all_state, - strict) - if load_all_state: - cls._restore_training_state(trainer, meta) - - if meta is not None: - _version = meta.get('modelscope') - if _version is not None and version.parse( - _version) < version.parse( - LoadCheckpointHook._TWO_PTH_FILE_VERSION): - trainer.logger.warning( - 'The unique pth file is split into a model file and ' - f'a trainer file since version {LoadCheckpointHook._TWO_PTH_FILE_VERSION},' - 'consider re-training your model or ' - 'using a converting script to split the single pth file into two.' - ) - trainer.logger.info( - f'Checkpoint {filename} saving time: {meta.get("time")}, modelscope version: {_version}' - ) - return meta - - @staticmethod - def load_trainer_state(trainer, train_state_file, load_all_state): - """Load trainer state file. - """ - - optimizer = getattr(trainer, 'optimizer', - None) if load_all_state else None - lr_scheduler = getattr(trainer, 'lr_scheduler', - None) if load_all_state else None - return load_checkpoint(train_state_file, None, optimizer, lr_scheduler) - - def load_model_state(self, trainer, model_file, strict): - """Load model state file. - """ - return load_checkpoint(model_file, - trainer.unwrap_module(trainer.model), None, - None) - - @Hook.overload_func(name='LoadCheckpointHook.load_checkpoints') - def load_checkpoints(self, checkpoint_path_prefix, trainer, load_all_state, - strict): - """Load checkpoint files of trainer state and model state. - - This is a strategic function which can be registered by other hook's function. - - Args: - checkpoint_path_prefix(str): The checkpoint dir with prefix or a model state file. - Example: '/tmp/test/epoch_0' or '/tmp/test/epoch_0.pth' - trainer(`EpochBasedTrainer`): The trainer instance. - load_all_state(`boolean`): Load all states (else load only module states). - strict(`boolean`): If strict, any unmatched keys will cause an error. - - Returns: - The meta info in json. - """ - _model_file, _train_state_file = _get_state_file_name( - checkpoint_path_prefix) - meta = {} - if os.path.isfile(_train_state_file): - meta = self.load_trainer_state(trainer, _train_state_file, - load_all_state) - else: - print(f'No trainer state file {_train_state_file} found, skip.') - self.load_model_state(trainer, _model_file, strict) - return meta - - -def _get_state_file_name(checkpoint_path_prefix): - """Get the default file name for state files. - - If the input is a checkpoint dir with prefix, this function will append suffix for both checkpoint files. - If the input is an absolute file name, this function will return it as the model file name, and append - suffix for the trainer file name. - - NOTE: a best checkpoint filename with float or int metric value inside - will not be judged as having a extension file name. like: '/tmp/test/epoch_0_accuracy0.85' - - Args: - checkpoint_path_prefix(`str`): The checkpoint dir with prefix or a model state file with extension file name. - like: '/tmp/test/epoch_0' - - Returns: - A tuple of model state file name and trainer state file name. - """ - base, ext = os.path.splitext(checkpoint_path_prefix) - if len(ext) == 0 or re.match(r'^\d+$', ext[1:]): - return checkpoint_path_prefix + CheckpointHook.MODEL_STATE_SUFFIX, \ - checkpoint_path_prefix + CheckpointHook.TRAINER_STATE_SUFFIX - else: - return checkpoint_path_prefix, base + CheckpointHook.TRAINER_STATE_SUFFIX.split( - '.')[0] + '.' + ext[1:] diff --git a/modelscope/trainers/hooks/compression/sparsity_hook.py b/modelscope/trainers/hooks/compression/sparsity_hook.py index 993488d8..e71c269a 100644 --- a/modelscope/trainers/hooks/compression/sparsity_hook.py +++ b/modelscope/trainers/hooks/compression/sparsity_hook.py @@ -1,7 +1,6 @@ # Copyright (c) Alibaba, Inc. and its affiliates. import os -from modelscope import __version__ from modelscope.metainfo import Hooks from modelscope.trainers.hooks.builder import HOOKS from modelscope.trainers.hooks.hook import Hook diff --git a/modelscope/trainers/hooks/distributed/__init__.py b/modelscope/trainers/hooks/distributed/__init__.py new file mode 100644 index 00000000..17b9c429 --- /dev/null +++ b/modelscope/trainers/hooks/distributed/__init__.py @@ -0,0 +1,3 @@ +from .ddp_hook import DDPHook +from .deepspeed_hook import DeepspeedHook +from .megatron_hook import MegatronHook diff --git a/modelscope/trainers/hooks/ddp_hook.py b/modelscope/trainers/hooks/distributed/ddp_hook.py similarity index 89% rename from modelscope/trainers/hooks/ddp_hook.py rename to modelscope/trainers/hooks/distributed/ddp_hook.py index eaae2d89..2bdbe939 100644 --- a/modelscope/trainers/hooks/ddp_hook.py +++ b/modelscope/trainers/hooks/distributed/ddp_hook.py @@ -1,11 +1,11 @@ # Copyright (c) Alibaba, Inc. and its affiliates. from modelscope.metainfo import Hooks +from modelscope.trainers.hooks.builder import HOOKS +from modelscope.trainers.hooks.hook import Hook +from modelscope.trainers.hooks.priority import Priority from modelscope.utils.constant import DistributedParallelType from modelscope.utils.device import create_device from modelscope.utils.torch_utils import get_local_rank, init_dist -from .builder import HOOKS -from .hook import Hook -from .priority import Priority @HOOKS.register_module(module_name=Hooks.DDPHook) diff --git a/modelscope/trainers/hooks/deepspeed_hook.py b/modelscope/trainers/hooks/distributed/deepspeed_hook.py similarity index 64% rename from modelscope/trainers/hooks/deepspeed_hook.py rename to modelscope/trainers/hooks/distributed/deepspeed_hook.py index a34b3f6f..7dddc5d9 100644 --- a/modelscope/trainers/hooks/deepspeed_hook.py +++ b/modelscope/trainers/hooks/distributed/deepspeed_hook.py @@ -8,72 +8,48 @@ from deepspeed import DeepSpeedEngine from megatron_util import mpu, print_rank_0 from modelscope.metainfo import Hooks +from modelscope.trainers.hooks import LoadCheckpointHook from modelscope.trainers.hooks.builder import HOOKS +from modelscope.trainers.hooks.checkpoint.checkpoint_hook import ( + BestCkptSaverHook, CheckpointHook) from modelscope.trainers.hooks.hook import Hook from modelscope.trainers.hooks.priority import Priority from modelscope.utils.checkpoint import save_checkpoint from modelscope.utils.logger import get_logger -from .checkpoint_hook import CheckpointHook, LoadCheckpointHook -from .megatron_hook import MegatronHook +from ..checkpoint.checkpoint_processor import CheckpointProcessor +from ..lr_scheduler_hook import LrSchedulerProcessor +from ..optimizer.base import OptimizerHook, OptimizerProcessor -@HOOKS.register_module(module_name=Hooks.DeepspeedHook) -class DeepspeedHook(MegatronHook): - PRIORITY = Priority.VERY_HIGH +class DeepspeedProcessor(CheckpointProcessor, LrSchedulerProcessor, + OptimizerProcessor): - def __init__(self, - deepspeed_activation_checkpointing=True, - save_zero_checkpoint=False, - with_mpu=True): - self.save_zero_checkpoint = save_zero_checkpoint - self.deepspeed_activation_checkpointing = deepspeed_activation_checkpointing - # TODO without mpu - self.with_mpu = with_mpu - assert with_mpu, 'DeepspeedHook now is only for mpu models.' + _BIN_FILE_DIR = 'model' - def register_strategy(self): - Hook.overload(name='OptimizerHook.backward', function=self.backward) - Hook.overload( - name='OptimizerHook.initialize_optimizer', function=self.idle) - Hook.overload(name='LrSchedulerHook.step', function=self.idle) - Hook.overload( - name='CheckpointHook.save_checkpoints', - function=self.save_checkpoints) - Hook.overload( - name='LoadCheckpointHook.load_checkpoints', - function=self.load_checkpoints) - Hook.overload( - name='CheckpointHook.remove_checkpoints', - function=self.remove_checkpoints) - Hook.overload( - name='CheckpointHook.prepare_output', function=self.prepare_output) - if self.with_mpu: - Hook.overload( - name='CheckpointHook.should_save_on_rank', - function=self.should_save_on_rank) + def rank_name(self): + # TODO + try: + tp_world_size = mpu.get_tensor_model_parallel_world_size() + if tp_world_size == 1: + return '' + mp_rank = mpu.get_tensor_model_parallel_rank() + return '_mp_rank_{:02d}'.format(mp_rank) + except (ImportError, AssertionError): + return '' - def backward(self, trainer, loss_keys, cumulative_iters, grad_clip): - # assert cumulative_iters == 1, 'DeepSpeed only support cumulative_iters=1' - # The `trainer.model` here is actually a deepspeed engine object. - # backward step - for k in loss_keys: - loss = trainer.train_outputs[k] - trainer.model.backward(loss) - - # update parameters - trainer.model.step() - - def idle(self, *args, **kwargs): - pass + def get_bin_file(self): + mp_rank = mpu.get_tensor_model_parallel_rank() + rank = '{:02d}'.format(mp_rank) + return f'mp_rank_{rank}_model_states.pt' def save_checkpoints(self, trainer, checkpoint_path_prefix, - output_sub_dir, + output_dir, meta=None): model = trainer.unwrap_module(trainer.model) _train_state_file = checkpoint_path_prefix + self.rank_name( - ) + CheckpointHook.TRAINER_STATE_SUFFIX + ) + CheckpointProcessor.TRAINER_STATE_SUFFIX # Save pth file without model state_dict save_checkpoint( model, _train_state_file, None, None, meta=meta, with_model=False) @@ -84,16 +60,22 @@ class DeepspeedHook(MegatronHook): bin_file = self.get_bin_file() src_file = os.path.join(checkpoint_path_prefix, bin_file) - dest_file = os.path.join(save_dir, output_sub_dir, self._BIN_FILE_DIR, - bin_file) + dest_file = os.path.join(output_dir, self._BIN_FILE_DIR, bin_file) if os.path.isfile(dest_file): os.unlink(dest_file) - os.link(src_file, dest_file) + try: + os.link(src_file, dest_file) + except OSError as e: + get_logger().error( + f'Link {src_file} to {dest_file} error: {e}, ' + 'changing to copy the bin file, this may case more space usage.' + ) + shutil.copyfile(src_file, dest_file) def remove_checkpoints(self, trainer, checkpoint_path_prefix): _train_state_file = checkpoint_path_prefix + self.rank_name( - ) + CheckpointHook.TRAINER_STATE_SUFFIX + ) + CheckpointProcessor.TRAINER_STATE_SUFFIX if os.path.isfile(_train_state_file): os.remove(_train_state_file) @@ -107,10 +89,10 @@ class DeepspeedHook(MegatronHook): meta = {} _train_state_file = checkpoint_path_prefix + self.rank_name( - ) + CheckpointHook.TRAINER_STATE_SUFFIX + ) + CheckpointProcessor.TRAINER_STATE_SUFFIX if os.path.isfile(_train_state_file): - meta = LoadCheckpointHook.load_trainer_state( - trainer, _train_state_file, load_all_state) + meta = self.load_trainer_state(trainer, _train_state_file, + load_all_state) if isinstance(trainer.model, DeepSpeedEngine): # DeepSpeedEngine is initialized @@ -138,6 +120,57 @@ class DeepspeedHook(MegatronHook): checkpoint, strict=strict) return meta + def backward(self, trainer, loss_keys, cumulative_iters, grad_clip): + # assert cumulative_iters == 1, 'DeepSpeed only support cumulative_iters=1' + # The `trainer.model` here is actually a deepspeed engine object. + # backward step + for k in loss_keys: + loss = trainer.train_outputs[k] + trainer.model.backward(loss) + + # update parameters + trainer.model.step() + + def initialize_optimizer(self, trainer): + pass + + def step(self, trainer): + pass + + +@HOOKS.register_module(module_name=Hooks.DeepspeedHook) +class DeepspeedHook(Hook): + PRIORITY = Priority.VERY_HIGH + + def __init__(self, + deepspeed_activation_checkpointing=True, + save_zero_checkpoint=False, + with_mpu=True): + self.save_zero_checkpoint = save_zero_checkpoint + self.deepspeed_activation_checkpointing = deepspeed_activation_checkpointing + # TODO without mpu + self.with_mpu = with_mpu + assert with_mpu, 'DeepspeedHook now is only for mpu models.' + + def register_processor(self, trainer): + processor = DeepspeedProcessor() + optimizer_hook = trainer.get_hook(OptimizerHook) + if len(optimizer_hook) > 0 and not isinstance( + optimizer_hook[0].processor, DeepspeedProcessor): + optimizer_hook[0].set_processor(processor) + ckpt_hook = trainer.get_hook(CheckpointHook) + if len(ckpt_hook) > 0 and not isinstance(ckpt_hook[0].processor, + DeepspeedProcessor): + ckpt_hook[0].set_processor(processor) + best_ckpt_hook = trainer.get_hook(BestCkptSaverHook) + if len(best_ckpt_hook) > 0 and not isinstance( + best_ckpt_hook[0].processor, DeepspeedProcessor): + best_ckpt_hook[0].set_processor(processor) + load_ckpt_hook = trainer.get_hook(LoadCheckpointHook) + if len(load_ckpt_hook) > 0 and not isinstance( + load_ckpt_hook[0].processor, DeepspeedProcessor): + load_ckpt_hook[0].set_processor(processor) + def before_val(self, trainer): pass diff --git a/modelscope/trainers/hooks/megatron_hook.py b/modelscope/trainers/hooks/distributed/megatron_hook.py similarity index 70% rename from modelscope/trainers/hooks/megatron_hook.py rename to modelscope/trainers/hooks/distributed/megatron_hook.py index f01288de..c4aeaf19 100644 --- a/modelscope/trainers/hooks/megatron_hook.py +++ b/modelscope/trainers/hooks/distributed/megatron_hook.py @@ -1,19 +1,129 @@ import os -from copy import deepcopy +import shutil import torch from megatron_util import mpu from modelscope.metainfo import Hooks +from modelscope.trainers import EpochBasedTrainer from modelscope.trainers.hooks.builder import HOOKS +from modelscope.trainers.hooks.checkpoint.checkpoint_hook import ( + BestCkptSaverHook, CheckpointHook, CheckpointProcessor) +from modelscope.trainers.hooks.checkpoint.load_checkpoint_hook import \ + LoadCheckpointHook from modelscope.trainers.hooks.hook import Hook -from modelscope.trainers.parallel.builder import build_parallel from modelscope.utils.checkpoint import load_checkpoint, save_checkpoint from modelscope.utils.constant import DistributedParallelType from modelscope.utils.device import create_device +from modelscope.utils.logger import get_logger from modelscope.utils.megatron_utils import is_megatron_initialized from modelscope.utils.torch_utils import get_local_rank -from .checkpoint_hook import CheckpointHook, LoadCheckpointHook + + +class MpuProcessor(CheckpointProcessor): + + _BIN_FILE_DIR = 'model' + + def rank_name(self): + # TODO + try: + tp_world_size = mpu.get_tensor_model_parallel_world_size() + if tp_world_size == 1: + return '' + mp_rank = mpu.get_tensor_model_parallel_rank() + return '_mp_rank_{:02d}'.format(mp_rank) + except (ImportError, AssertionError): + return '' + + def get_bin_file(self): + mp_rank = mpu.get_tensor_model_parallel_rank() + rank = '{:02d}'.format(mp_rank) + return f'mp_rank_{rank}_model_states.pt' + + def should_save_on_rank(self, trainer): + # TODO + return (not torch.distributed.is_initialized() + ) or mpu.get_data_parallel_rank() == 0 + + def prepare_output(self, trainer, output_dir): + config = trainer.cfg + CheckpointProcessor.copy_files_and_dump_config(trainer, output_dir, + config, + self._BIN_FILE_DIR) + os.makedirs( + os.path.join(output_dir, self._BIN_FILE_DIR), exist_ok=True) + + def save_checkpoints(self, + trainer, + checkpoint_path_prefix, + output_dir, + meta=None): + model = trainer.unwrap_module(trainer.model) + _train_state_file = checkpoint_path_prefix + self.rank_name( + ) + CheckpointProcessor.TRAINER_STATE_SUFFIX + # Save pth file without model state_dict + save_checkpoint( + model, + _train_state_file, + trainer.optimizer, + trainer.lr_scheduler, + meta=meta, + with_model=False) + + save_dir = os.path.dirname(checkpoint_path_prefix) + prefix = os.path.basename(checkpoint_path_prefix) + bin_file = self.get_bin_file() + prefix_bin_file = os.path.join(save_dir, prefix + '_' + bin_file) + save_checkpoint(model, prefix_bin_file, with_meta=False) + + src_file = prefix_bin_file + dest_file = os.path.join(output_dir, self._BIN_FILE_DIR, bin_file) + if os.path.isfile(dest_file): + os.unlink(dest_file) + + try: + os.link(src_file, dest_file) + except OSError as e: + get_logger().error( + f'Link {src_file} to {dest_file} error: {e}, ' + 'changing to copy the bin file, this may case more space usage.' + ) + shutil.copyfile(src_file, dest_file) + + def remove_checkpoints(self, trainer, checkpoint_path_prefix): + _train_state_file = checkpoint_path_prefix + self.rank_name( + ) + CheckpointProcessor.TRAINER_STATE_SUFFIX + if os.path.isfile(_train_state_file): + os.remove(_train_state_file) + + save_dir = os.path.dirname(checkpoint_path_prefix) + prefix = os.path.basename(checkpoint_path_prefix) + bin_file = self.get_bin_file() + absolute_file = os.path.join(save_dir, prefix + '_' + bin_file) + if os.path.isfile(absolute_file): + os.remove(absolute_file) + + def load_checkpoints(self, checkpoint_path_prefix, trainer, load_all_state, + strict): + model = trainer.unwrap_module(trainer.model) + if os.path.isdir(checkpoint_path_prefix): + save_dir = checkpoint_path_prefix + bin_file = self.get_bin_file() + model_file = os.path.join(save_dir, bin_file) + load_checkpoint(model_file, model, None, None) + else: + _train_state_file = checkpoint_path_prefix + self.rank_name( + ) + CheckpointProcessor.TRAINER_STATE_SUFFIX + meta = LoadCheckpointHook.load_trainer_state( + trainer, _train_state_file, load_all_state) + + save_dir = os.path.dirname(checkpoint_path_prefix) + prefix = os.path.basename(checkpoint_path_prefix) + bin_file = self.get_bin_file() + + model_file = os.path.join(save_dir, prefix + '_' + bin_file) + load_checkpoint(model_file, model, None, None) + return meta @HOOKS.register_module(module_name=Hooks.MegatronHook) @@ -24,21 +134,20 @@ class MegatronHook(Hook): def __init__(self): self.wrapped = False - def register_strategy(self): - Hook.overload( - name='CheckpointHook.should_save_on_rank', - function=self.should_save_on_rank) - Hook.overload( - name='CheckpointHook.save_checkpoints', - function=self.save_checkpoints) - Hook.overload( - name='LoadCheckpointHook.load_checkpoints', - function=self.load_checkpoints) - Hook.overload( - name='CheckpointHook.remove_checkpoints', - function=self.remove_checkpoints) - Hook.overload( - name='CheckpointHook.prepare_output', function=self.prepare_output) + def register_processor(self, trainer: EpochBasedTrainer): + processor = MpuProcessor() + ckpt_hook = trainer.get_hook(CheckpointHook) + if len(ckpt_hook) > 0 and not isinstance(ckpt_hook[0].processor, + MpuProcessor): + ckpt_hook[0].set_processor(processor) + best_ckpt_hook = trainer.get_hook(BestCkptSaverHook) + if len(best_ckpt_hook) > 0 and not isinstance( + best_ckpt_hook[0].processor, MpuProcessor): + best_ckpt_hook[0].set_processor(processor) + load_ckpt_hook = trainer.get_hook(LoadCheckpointHook) + if len(load_ckpt_hook) > 0 and not isinstance( + load_ckpt_hook[0].processor, MpuProcessor): + load_ckpt_hook[0].set_processor(processor) def after_init(self, trainer): assert is_megatron_initialized() @@ -63,97 +172,3 @@ class MegatronHook(Hook): if not self.wrapped: trainer.model = trainer.to_parallel(trainer.model) self.wrapped = True - - def should_save_on_rank(self, trainer): - # TODO - return (not torch.distributed.is_initialized() - ) or mpu.get_data_parallel_rank() == 0 - - def rank_name(self): - # TODO - try: - tp_world_size = mpu.get_tensor_model_parallel_world_size() - if tp_world_size == 1: - return '' - mp_rank = mpu.get_tensor_model_parallel_rank() - return '_mp_rank_{:02d}'.format(mp_rank) - except (ImportError, AssertionError): - return '' - - def get_bin_file(self): - mp_rank = mpu.get_tensor_model_parallel_rank() - rank = '{:02d}'.format(mp_rank) - return f'mp_rank_{rank}_model_states.pt' - - def save_checkpoints(self, - trainer, - checkpoint_path_prefix, - output_sub_dir, - meta=None): - model = trainer.unwrap_module(trainer.model) - _train_state_file = checkpoint_path_prefix + self.rank_name( - ) + CheckpointHook.TRAINER_STATE_SUFFIX - # Save pth file without model state_dict - save_checkpoint( - model, - _train_state_file, - trainer.optimizer, - trainer.lr_scheduler, - meta=meta, - with_model=False) - - save_dir = os.path.dirname(checkpoint_path_prefix) - prefix = os.path.basename(checkpoint_path_prefix) - bin_file = self.get_bin_file() - prefix_bin_file = os.path.join(save_dir, prefix + '_' + bin_file) - save_checkpoint(model, prefix_bin_file, with_meta=False) - - src_file = prefix_bin_file - dest_file = os.path.join(save_dir, output_sub_dir, self._BIN_FILE_DIR, - bin_file) - if os.path.isfile(dest_file): - os.unlink(dest_file) - - os.link(src_file, dest_file) - - def remove_checkpoints(self, trainer, checkpoint_path_prefix): - _train_state_file = checkpoint_path_prefix + self.rank_name( - ) + CheckpointHook.TRAINER_STATE_SUFFIX - if os.path.isfile(_train_state_file): - os.remove(_train_state_file) - - save_dir = os.path.dirname(checkpoint_path_prefix) - prefix = os.path.basename(checkpoint_path_prefix) - bin_file = self.get_bin_file() - absolute_file = os.path.join(save_dir, prefix + '_' + bin_file) - if os.path.isfile(absolute_file): - os.remove(absolute_file) - - def load_checkpoints(self, checkpoint_path_prefix, trainer, load_all_state, - strict): - model = trainer.unwrap_module(trainer.model) - if os.path.isdir(checkpoint_path_prefix): - save_dir = checkpoint_path_prefix - bin_file = self.get_bin_file() - model_file = os.path.join(save_dir, bin_file) - load_checkpoint(model_file, model, None, None) - else: - _train_state_file = checkpoint_path_prefix + self.rank_name( - ) + CheckpointHook.TRAINER_STATE_SUFFIX - meta = LoadCheckpointHook.load_trainer_state( - trainer, _train_state_file, load_all_state) - - save_dir = os.path.dirname(checkpoint_path_prefix) - prefix = os.path.basename(checkpoint_path_prefix) - bin_file = self.get_bin_file() - - model_file = os.path.join(save_dir, prefix + '_' + bin_file) - load_checkpoint(model_file, model, None, None) - return meta - - def prepare_output(self, trainer, output_dir): - config = trainer.cfg - CheckpointHook.copy_files_and_dump_config(trainer, output_dir, config, - self._BIN_FILE_DIR) - os.makedirs( - os.path.join(output_dir, self._BIN_FILE_DIR), exist_ok=True) diff --git a/modelscope/trainers/hooks/early_stop_hook.py b/modelscope/trainers/hooks/early_stop_hook.py index b15e8e5a..7aba69a4 100644 --- a/modelscope/trainers/hooks/early_stop_hook.py +++ b/modelscope/trainers/hooks/early_stop_hook.py @@ -9,6 +9,12 @@ from .hook import Hook from .priority import Priority +class EarlyStopStrategy: + by_epoch = 'by_epoch' + by_step = 'by_step' + no = 'no' + + @HOOKS.register_module(module_name=Hooks.EarlyStopHook) class EarlyStopHook(Hook): """Early stop when a specific metric stops improving. @@ -16,14 +22,13 @@ class EarlyStopHook(Hook): Args: metric_key (str): Metric key to be monitored. rule (str): Comparison rule for best score. Support "max" and "min". - If rule is "max", the training will stop when `metric_key` has stopped increaing. + If rule is "max", the training will stop when `metric_key` has stopped increasing. If rule is "min", the training will stop when `metric_key` has stopped decreasing. patience (int): Trainer will stop if the monitored metric did not improve for the last `patience` times. - min_delta (float): Minimum change in the monitored metric to quailfy as an improvement. + min_delta (float): Minimum change in the monitored metric to qualify as an improvement. check_finite (bool): If true, stops training when the metric becomes NaN or infinite. - by_epoch (int): Saving checkpoints by epoch or by iteration. - interval (int): The frequency to trigger early stop check. If `by_epoch=True`, - it means the number of epochs, else means the number of iterations. + early_stop_strategy (str): The strategy to early stop, can be by_epoch/by_step/none + interval (int): The frequency to trigger early stop check, by epoch or step. """ PRIORITY = Priority.VERY_LOW @@ -35,14 +40,19 @@ class EarlyStopHook(Hook): patience: int = 3, min_delta: float = 0.0, check_finite: bool = True, - by_epoch: bool = True, - interval: int = 1): + early_stop_strategy: str = EarlyStopStrategy.by_epoch, + interval: int = 1, + **kwargs): self.metric_key = metric_key self.rule = rule self.patience = patience self.min_delta = min_delta self.check_finite = check_finite - self.by_epoch = by_epoch + if 'by_epoch' in kwargs: + self.early_stop_strategy = EarlyStopStrategy.by_epoch if kwargs[ + 'by_epoch'] else EarlyStopStrategy.by_step + else: + self.early_stop_strategy = early_stop_strategy self.interval = interval self.wait_count = 0 @@ -89,7 +99,7 @@ class EarlyStopHook(Hook): trainer._stop_training = True def after_train_epoch(self, trainer): - if not self.by_epoch: + if self.early_stop_strategy != EarlyStopStrategy.by_epoch: return if not self.every_n_epochs(trainer, self.interval): @@ -99,7 +109,7 @@ class EarlyStopHook(Hook): self._stop_training(trainer) def after_train_iter(self, trainer): - if self.by_epoch: + if self.early_stop_strategy != EarlyStopStrategy.by_step: return if not self.every_n_iters(trainer, self.interval): diff --git a/modelscope/trainers/hooks/evaluation_hook.py b/modelscope/trainers/hooks/evaluation_hook.py index 80c8c31a..c29a6d6a 100644 --- a/modelscope/trainers/hooks/evaluation_hook.py +++ b/modelscope/trainers/hooks/evaluation_hook.py @@ -1,11 +1,18 @@ # Copyright (c) Alibaba, Inc. and its affiliates. from collections import OrderedDict +from typing import Optional from modelscope.metainfo import Hooks from .builder import HOOKS from .hook import Hook +class EvaluationStrategy: + by_epoch = 'by_epoch' + by_step = 'by_step' + no = 'no' + + @HOOKS.register_module(module_name=Hooks.EvaluationHook) class EvaluationHook(Hook): """ @@ -18,21 +25,34 @@ class EvaluationHook(Hook): Default: None, validate every interval epochs/iterations from scratch. """ - def __init__(self, interval=1, by_epoch=True, start_idx=None): + def __init__(self, + interval: Optional[int] = 1, + eval_strategy: Optional[str] = EvaluationStrategy.by_epoch, + start_idx: Optional[int] = None, + **kwargs): assert interval > 0, 'interval must be a positive number' self.interval = interval self.start_idx = start_idx - self.by_epoch = by_epoch + self.last_eval_tag = (None, None) + if 'by_epoch' in kwargs: + self.eval_strategy = EvaluationStrategy.by_epoch if kwargs[ + 'by_epoch'] else EvaluationStrategy.by_step + else: + self.eval_strategy = eval_strategy def after_train_iter(self, trainer): """Called after every training iter to evaluate the results.""" - if not self.by_epoch and self._should_evaluate(trainer): + if self.eval_strategy == EvaluationStrategy.by_step and self._should_evaluate( + trainer): self.do_evaluate(trainer) + self.last_eval_tag = ('iter', trainer.iter) def after_train_epoch(self, trainer): """Called after every training epoch to evaluate the results.""" - if self.by_epoch and self._should_evaluate(trainer): + if self.eval_strategy == EvaluationStrategy.by_epoch and self._should_evaluate( + trainer): self.do_evaluate(trainer) + self.last_eval_tag = ('epoch', trainer.epoch) def add_visualization_info(self, trainer, results): if trainer.visualization_buffer.output.get('eval_results', @@ -64,7 +84,7 @@ class EvaluationHook(Hook): Returns: bool: The flag indicating whether to perform evaluation. """ - if self.by_epoch: + if self.eval_strategy == EvaluationStrategy.by_epoch: current = trainer.epoch check_time = self.every_n_epochs else: diff --git a/modelscope/trainers/hooks/hook.py b/modelscope/trainers/hooks/hook.py index 70e06fbd..93ea8541 100644 --- a/modelscope/trainers/hooks/hook.py +++ b/modelscope/trainers/hooks/hook.py @@ -22,9 +22,6 @@ class Hook: PRIORITY = Priority.NORMAL - # The strategic function dict. - _strategies = dict() - def after_init(self, trainer): """ Will be called at the end of the trainer's `__init__` method @@ -201,42 +198,48 @@ class Hook: """ self.after_iter(trainer) - def every_n_epochs(self, trainer, n): + @staticmethod + def every_n_epochs(trainer, n): """ Whether to reach every ``n`` epochs Returns: bool """ return (trainer.epoch + 1) % n == 0 if n > 0 else False - def every_n_inner_iters(self, runner, n): + @staticmethod + def every_n_inner_iters(runner, n): """ Whether to reach every ``n`` iterations at every epoch Returns: bool """ return (runner.inner_iter + 1) % n == 0 if n > 0 else False - def every_n_iters(self, trainer, n): + @staticmethod + def every_n_iters(trainer, n): """ Whether to reach every ``n`` iterations Returns: bool """ return (trainer.iter + 1) % n == 0 if n > 0 else False - def end_of_epoch(self, trainer): + @staticmethod + def end_of_epoch(trainer): """ Whether to reach the end of every epoch Returns: bool """ return trainer.inner_iter + 1 == trainer.iters_per_epoch - def is_last_epoch(self, trainer): + @staticmethod + def is_last_epoch(trainer): """ Whether to reach the last epoch Returns: bool """ return trainer.epoch + 1 == trainer.max_epochs - def is_last_iter(self, trainer): + @staticmethod + def is_last_iter(trainer): """ Whether to reach the last iteration in the entire training process Returns: bool @@ -256,54 +259,3 @@ class Hook: def load_state_dict(self, state_dict): pass - - @staticmethod - def clear_strategies(): - Hook._strategies.clear() - - @staticmethod - def overload(function, name=None): - """Register a function to a strategic function. - - Args: - function(`method` or `Callable`): The function instance. - name(`str`): The name of the strategic function, which specifies by the method `consume` - """ - - _name = name or function.__name__ - if _name not in Hook._strategies: - Hook._strategies[_name] = [] - - Hook._strategies[_name].append(function) - - @staticmethod - def overload_func(name=None): - """Declare a function as a strategic function, which can be replaced by some other functions. - - This function should be used in annotations. - - Args: - name(str): The strategic function name. - """ - - def _register(function): - - @wraps(function) - def _call(*args, **kwargs): - _name = name or function.__name__ - producers = Hook._strategies.get(_name, []) - - if len(producers) == 0: - return function(*args, **kwargs) - else: - if len(producers) > 1: - raise ValueError( - f'Multiple functions registered to {_name}, ' - f'here is the list: {producers}') - if isinstance(args[0], Hook): - args = args[1:] - return producers[0](*args, **kwargs) - - return _call - - return _register diff --git a/modelscope/trainers/hooks/lr_scheduler_hook.py b/modelscope/trainers/hooks/lr_scheduler_hook.py index 28ce250c..51a8e858 100644 --- a/modelscope/trainers/hooks/lr_scheduler_hook.py +++ b/modelscope/trainers/hooks/lr_scheduler_hook.py @@ -1,4 +1,5 @@ # Copyright (c) Alibaba, Inc. and its affiliates. + from modelscope.metainfo import Hooks from modelscope.trainers.lrscheduler.builder import build_lr_scheduler from modelscope.utils.constant import LogKeys @@ -9,6 +10,42 @@ from .hook import Hook from .priority import Priority +class LrSchedulerProcessor: + + def __init__(self): + self.lr_strategy = None + self.warmup_lr_scheduler = None + + def set_lr_strategy(self, lr_strategy): + self.lr_strategy = lr_strategy + + def set_warmup_lr_scheduler(self, warmup_lr_scheduler): + self.warmup_lr_scheduler = warmup_lr_scheduler + + def initialize_lr_scheduler(self, trainer): + """Initialize the lr scheduler. + + This is a strategic function which can be registered by other hook's function. + """ + pass + + def step(self, trainer): + """Do lr scheduler's step. + + This is a strategic function which can be registered by other hook's function. + """ + if self.warmup_lr_scheduler is not None: + self.warmup_lr_scheduler.step() + else: + trainer.lr_scheduler.step() + + +class LrStrategy: + by_epoch = 'by_epoch' + by_step = 'by_step' + no = 'no' + + @HOOKS.register_module(module_name=Hooks.LrSchedulerHook) class LrSchedulerHook(Hook): """Lr scheduler. @@ -19,38 +56,33 @@ class LrSchedulerHook(Hook): """ PRIORITY = Priority.LOW - def __init__(self, by_epoch=True, warmup=None, **kwargs) -> None: + def __init__(self, + lr_strategy=LrStrategy.by_epoch, + warmup=None, + **kwargs) -> None: super().__init__() - self.by_epoch = by_epoch + if 'by_epoch' in kwargs: + self.lr_strategy = LrStrategy.by_epoch if kwargs[ + 'by_epoch'] else LrStrategy.by_step + else: + self.lr_strategy = lr_strategy self.warmup = warmup self.warmup_lr_scheduler = None + self.processor = LrSchedulerProcessor() + + def set_processor(self, processor): + self.processor = processor def before_run(self, trainer): - self.initialize_lr_scheduler(trainer) + self.processor.set_lr_strategy(self.lr_strategy) if self.warmup is not None: assert isinstance(self.warmup, dict) and 'type' in self.warmup self.warmup_lr_scheduler = build_lr_scheduler( cfg=self.warmup, default_args={'base_scheduler': trainer.lr_scheduler}) + self.processor.set_warmup_lr_scheduler(self.warmup_lr_scheduler) - @Hook.overload_func(name='LrSchedulerHook.initialize_lr_scheduler') - def initialize_lr_scheduler(self, trainer): - """Initialize the lr scheduler. - - This is a strategic function which can be registered by other hook's function. - """ - pass - - @Hook.overload_func(name='LrSchedulerHook.step') - def step(self, trainer): - """Do lr scheduler's step. - - This is a strategic function which can be registered by other hook's function. - """ - if self.warmup_lr_scheduler is not None: - self.warmup_lr_scheduler.step() - else: - trainer.lr_scheduler.step() + self.processor.initialize_lr_scheduler(trainer) def get_current_lr(self, trainer): import torch @@ -67,17 +99,17 @@ class LrSchedulerHook(Hook): return lr def after_train_iter(self, trainer): - if not self.by_epoch and trainer.iter >= getattr( + if self.lr_strategy == LrStrategy.by_step and trainer.iter >= getattr( trainer, 'cumulative_iters', 1) - 1: - self.step(trainer) + self.processor.step(trainer) trainer.log_buffer.output[LogKeys.LR] = self._get_log_lr(trainer) def before_train_epoch(self, trainer): trainer.log_buffer.output[LogKeys.LR] = self._get_log_lr(trainer) def after_train_epoch(self, trainer): - if self.by_epoch: - self.step(trainer) + if self.lr_strategy == LrStrategy.by_epoch: + self.processor.step(trainer) def _get_log_lr(self, trainer): cur_lr = self.get_current_lr(trainer) @@ -94,6 +126,29 @@ class LrSchedulerHook(Hook): return lr +class PlateauLrSchedulerProcessor(LrSchedulerProcessor): + + def __init__(self, metric_key): + super().__init__() + self.metric_key = metric_key + + def step(self, trainer): + # adapt to evaluation interval is greater than 1 + if trainer.metric_values is None: + if is_master(): + print( + f'Current epoch {trainer.epoch} has no evaluation metric values, skip lr_scheduler.step() !' + ) + return + + metrics = trainer.metric_values[self.metric_key] + if self.lr_strategy == LrStrategy.by_epoch: + if self.warmup_lr_scheduler is not None: + self.warmup_lr_scheduler.step(metrics=metrics) + else: + trainer.lr_scheduler.step(metrics=metrics) + + @HOOKS.register_module(module_name=Hooks.PlateauLrSchedulerHook) class PlateauLrSchedulerHook(Hook): """Lr scheduler hook for `ReduceLROnPlateau`. @@ -105,10 +160,16 @@ class PlateauLrSchedulerHook(Hook): PRIORITY = Priority.LOW # should be after EvaluationHook def __init__(self, metric_key, **kwargs): + super().__init__() self.metric_key = metric_key - def register_strategy(self): - Hook.overload(name='LrSchedulerHook.step', function=self.step) + def register_processor(self, trainer): + lr_scheduler_hook = trainer.get_hook(LrSchedulerHook) + if len(lr_scheduler_hook) > 0 and type( + lr_scheduler_hook[0].processor) in (type(None), + LrSchedulerProcessor): + lr_scheduler_hook[0].set_processor( + PlateauLrSchedulerProcessor(self.metric_key)) def before_run(self, trainer): if not hasattr(trainer, 'logger'): @@ -116,23 +177,6 @@ class PlateauLrSchedulerHook(Hook): else: self.logger = trainer.logger - def step(self, trainer): - # adapt to evaluation intervel is greater than 1 - if trainer.metric_values is None: - if is_master(): - self.logger.warning( - f'Current epoch {trainer.epoch} has no evaluation metric values, skip lr_scheduler.step() !' - ) - return - - metrics = trainer.metric_values[self.metric_key] - lr_scheduler_hook = trainer.get_hook(LrSchedulerHook)[0] - if lr_scheduler_hook.by_epoch: - if lr_scheduler_hook.warmup_lr_scheduler is not None: - lr_scheduler_hook.warmup_lr_scheduler.step(metrics=metrics) - else: - trainer.lr_scheduler.step(metrics=metrics) - @HOOKS.register_module(module_name=Hooks.NoneLrSchedulerHook) class NoneLrSchedulerHook(LrSchedulerHook): diff --git a/modelscope/trainers/hooks/optimizer/apex_optimizer_hook.py b/modelscope/trainers/hooks/optimizer/apex_optimizer_hook.py index bd1034f3..3c874ccf 100644 --- a/modelscope/trainers/hooks/optimizer/apex_optimizer_hook.py +++ b/modelscope/trainers/hooks/optimizer/apex_optimizer_hook.py @@ -7,40 +7,14 @@ from packaging import version from modelscope.metainfo import Hooks from modelscope.trainers.hooks import Hook from modelscope.trainers.hooks.builder import HOOKS -from .base import OptimizerHook +from .base import OptimizerHook, OptimizerProcessor -@HOOKS.register_module(module_name=Hooks.ApexAMPOptimizerHook) -class ApexAMPOptimizerHook(Hook): - """ - Fp16 optimizer, if torch version is less than 1.6.0, - you must install apex (https://www.github.com/nvidia/apex) else use torch.cuda.amp by default +class ApexOptimizerProcessor(OptimizerProcessor): - Args: - opt_level (str): "O0" and "O3" are not true mixed precision, - but they are useful for establishing accuracy and speed baselines, respectively. - "O1" and "O2" are different implementations of mixed precision. - Try both, and see what gives the best speedup and accuracy for your model. - """ - - PRIORITY = OptimizerHook.PRIORITY - - def __init__(self, opt_level='O1', **kwargs): + def __init__(self, opt_level): self.opt_level = opt_level - try: - from apex import amp - except ImportError: - raise ValueError( - 'apex not installed, please install apex from https://www.github.com/nvidia/apex.' - ) - - def register_strategy(self): - Hook.overload( - name='OptimizerHook.initialize_optimizer', - function=self.initialize_optimizer) - Hook.overload(name='OptimizerHook.backward', function=self.backward) - def initialize_optimizer(self, trainer): from apex import amp @@ -68,10 +42,44 @@ class ApexAMPOptimizerHook(Hook): trainer.optimizer) as scaled_loss: scaled_loss.backward() - if self.every_n_iters(trainer, cumulative_iters): + if Hook.every_n_iters(trainer, cumulative_iters): if grad_clip is not None: - OptimizerHook.clip_grads(trainer.model.parameters(), - **grad_clip) + OptimizerProcessor.clip_grads(trainer.model.parameters(), + **grad_clip) trainer.optimizer.step() trainer.optimizer.zero_grad() + + +@HOOKS.register_module(module_name=Hooks.ApexAMPOptimizerHook) +class ApexAMPOptimizerHook(Hook): + """ + Fp16 optimizer, if torch version is less than 1.6.0, + you must install apex (https://www.github.com/nvidia/apex) else use torch.cuda.amp by default + + Args: + opt_level (str): "O0" and "O3" are not true mixed precision, + but they are useful for establishing accuracy and speed baselines, respectively. + "O1" and "O2" are different implementations of mixed precision. + Try both, and see what gives the best speedup and accuracy for your model. + """ + + PRIORITY = OptimizerHook.PRIORITY + + def __init__(self, opt_level='O1', **kwargs): + self.opt_level = opt_level + + try: + from apex import amp + except ImportError: + raise ValueError( + 'apex not installed, please install apex from https://www.github.com/nvidia/apex.' + ) + + def register_processor(self, trainer): + optimizer_hook = trainer.get_hook(OptimizerHook) + if len(optimizer_hook) > 0 and type( + optimizer_hook[0].processor) in (type(None), + OptimizerProcessor): + optimizer_hook[0].set_processor( + ApexOptimizerProcessor(self.opt_level)) diff --git a/modelscope/trainers/hooks/optimizer/base.py b/modelscope/trainers/hooks/optimizer/base.py index f0d62612..ca20720d 100644 --- a/modelscope/trainers/hooks/optimizer/base.py +++ b/modelscope/trainers/hooks/optimizer/base.py @@ -10,6 +10,48 @@ from modelscope.trainers.hooks.hook import Hook from modelscope.trainers.hooks.priority import Priority +class OptimizerProcessor: + + def initialize_optimizer(self, trainer): + """Initialize the optimizer. + + This is a strategic function which can be registered by other hook's function. + """ + trainer.optimizer.zero_grad() + + def before_forward(self, trainer): + pass + + def backward(self, trainer, loss_keys, cumulative_iters, grad_clip): + """Do module backward, optimizer's step and zero_grad and clip the grads. + + This is a strategic function which can be registered by other hook's function. + + Args: + trainer(`EpochBasedTrainer`): The trainer instance. + loss_keys(`list`): The list of loss keys. + cumulative_iters(`int`): The cumulative iters for gradients. + grad_clip(`dict`): The grad clipping options. + """ + for k in loss_keys: + trainer.train_outputs[k] /= cumulative_iters + trainer.train_outputs[k].backward() + + if Hook.every_n_iters(trainer, cumulative_iters): + if grad_clip is not None: + self.clip_grads(trainer.model.parameters(), **grad_clip) + + trainer.optimizer.step() + trainer.optimizer.zero_grad() + + @staticmethod + def clip_grads(params, **clip_args): + params = list( + filter(lambda p: p.requires_grad and p.grad is not None, params)) + if len(params) > 0: + return clip_grad.clip_grad_norm_(params, **clip_args) + + @HOOKS.register_module(module_name=Hooks.OptimizerHook) class OptimizerHook(Hook): """Optimizer hook @@ -36,52 +78,21 @@ class OptimizerHook(Hook): self.loss_keys = loss_keys self.cumulative_iters = cumulative_iters self.grad_clip = grad_clip + self.processor = OptimizerProcessor() - @staticmethod - def clip_grads(params, **clip_args): - params = list( - filter(lambda p: p.requires_grad and p.grad is not None, params)) - if len(params) > 0: - return clip_grad.clip_grad_norm_(params, **clip_args) - - @Hook.overload_func(name='OptimizerHook.initialize_optimizer') - def initialize_optimizer(self, trainer): - """Initialize the optimizer. - - This is a strategic function which can be registered by other hook's function. - """ - trainer.optimizer.zero_grad() + def set_processor(self, processor): + self.processor = processor def before_run(self, trainer): - self.initialize_optimizer(trainer) trainer.cumulative_iters = self.cumulative_iters + self.processor.initialize_optimizer(trainer) - @Hook.overload_func(name='OptimizerHook.backward') - def backward(self, trainer, loss_keys, cumulative_iters, grad_clip): - """Do module backward, optimizer's step and zero_grad and clip the grads. - - This is a strategic function which can be registered by other hook's function. - - Args: - trainer(`EpochBasedTrainer`): The trainer instance. - loss_keys(`list`): The list of loss keys. - cumulative_iters(`int`): The cumulative iters for gradients. - grad_clip(`dict`): The grad clipping options. - """ - for k in loss_keys: - trainer.train_outputs[k] /= cumulative_iters - trainer.train_outputs[k].backward() - - if self.every_n_iters(trainer, cumulative_iters): - if grad_clip is not None: - self.clip_grads(trainer.model.parameters(), **grad_clip) - - trainer.optimizer.step() - trainer.optimizer.zero_grad() + def before_train_iter(self, trainer): + self.processor.before_forward(trainer) def after_train_iter(self, trainer): - self.backward(trainer, self.loss_keys, self.cumulative_iters, - self.grad_clip) + self.processor.backward(trainer, self.loss_keys, self.cumulative_iters, + self.grad_clip) @HOOKS.register_module(module_name=Hooks.NoneOptimizerHook) diff --git a/modelscope/trainers/hooks/optimizer/torch_optimizer_hook.py b/modelscope/trainers/hooks/optimizer/torch_optimizer_hook.py index 1ab89720..fc7d2672 100644 --- a/modelscope/trainers/hooks/optimizer/torch_optimizer_hook.py +++ b/modelscope/trainers/hooks/optimizer/torch_optimizer_hook.py @@ -4,7 +4,45 @@ import logging from modelscope.metainfo import Hooks from modelscope.trainers.hooks import Hook from modelscope.trainers.hooks.builder import HOOKS -from .base import OptimizerHook +from .base import OptimizerHook, OptimizerProcessor + + +class TorchAMPOptimizerProcessor(OptimizerProcessor): + + def __init__(self, scaler, scale_update_param): + self.scaler = scaler + self.scale_update_param = scale_update_param + + def before_forward(self, trainer): + from torch.cuda import amp + setattr(self._model, 'forward', amp.autocast()(self._model.forward)) + + def initialize_optimizer(self, trainer): + logging.info('open fp16') + trainer.optimizer.zero_grad() + + model = trainer.unwrap_module(trainer.model) + self._ori_model_forward = model.forward + self._model = model + + def backward(self, trainer, loss_keys, cumulative_iters, grad_clip): + for k in loss_keys: + trainer.train_outputs[k] /= cumulative_iters + + for k in loss_keys: + self.scaler.scale(trainer.train_outputs[k]).backward() + + if Hook.every_n_iters(trainer, cumulative_iters): + self.scaler.unscale_(trainer.optimizer) + if grad_clip is not None: + OptimizerProcessor.clip_grads(trainer.model.parameters(), + **grad_clip) + + self.scaler.step(trainer.optimizer) + self.scaler.update(self.scale_update_param) + trainer.optimizer.zero_grad() + + setattr(self._model, 'forward', self._ori_model_forward) @HOOKS.register_module(module_name=Hooks.TorchAMPOptimizerHook) @@ -44,39 +82,11 @@ class TorchAMPOptimizerHook(Hook): '`loss_scale` type must be in [float, dict], but got {loss_scale}' ) - def register_strategy(self): - Hook.overload( - name='OptimizerHook.initialize_optimizer', - function=self.initialize_optimizer) - Hook.overload(name='OptimizerHook.backward', function=self.backward) - - def initialize_optimizer(self, trainer): - logging.info('open fp16') - trainer.optimizer.zero_grad() - - model = trainer.unwrap_module(trainer.model) - self._ori_model_forward = model.forward - self._model = model - - def before_train_iter(self, trainer): - from torch.cuda import amp - setattr(self._model, 'forward', amp.autocast()(self._model.forward)) - - def backward(self, trainer, loss_keys, cumulative_iters, grad_clip): - for k in loss_keys: - trainer.train_outputs[k] /= cumulative_iters - - for k in loss_keys: - self.scaler.scale(trainer.train_outputs[k]).backward() - - if self.every_n_iters(trainer, cumulative_iters): - self.scaler.unscale_(trainer.optimizer) - if grad_clip is not None: - OptimizerHook.clip_grads(trainer.model.parameters(), - **grad_clip) - - self.scaler.step(trainer.optimizer) - self.scaler.update(self._scale_update_param) - trainer.optimizer.zero_grad() - - setattr(self._model, 'forward', self._ori_model_forward) + def register_processor(self, trainer): + optimizer_hook = trainer.get_hook(OptimizerHook) + if len(optimizer_hook) > 0 and type( + optimizer_hook[0].processor) in (type(None), + OptimizerProcessor): + optimizer_hook[0].set_processor( + TorchAMPOptimizerProcessor(self.scaler, + self._scale_update_param)) diff --git a/modelscope/trainers/trainer.py b/modelscope/trainers/trainer.py index 683ff2f5..c980de04 100644 --- a/modelscope/trainers/trainer.py +++ b/modelscope/trainers/trainer.py @@ -11,7 +11,7 @@ import json import torch from torch import distributed as dist from torch import nn -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import DataLoader, Dataset, Sampler from torch.utils.data.dataloader import default_collate from torch.utils.data.distributed import DistributedSampler @@ -88,7 +88,7 @@ class EpochBasedTrainer(BaseTrainer): compile_options (dict, optional): The compile options if compile=True, default None to use the default params of 'TorchModel.compile'. efficient_tuners (dict, optional): The tuners to use to train the model - + samplers: (:obj:`Sampler` or `Dict[Sampler]`, *optional*): samplers used in the train/eval DataLoader. Examples of cfg_modify_fn: >>> def cfg_modify_fn(cfg): >>> cfg.preprocessor.first_sequence= 'text1' @@ -114,6 +114,7 @@ class EpochBasedTrainer(BaseTrainer): model_revision: Optional[str] = DEFAULT_MODEL_REVISION, seed: int = 42, callbacks: Optional[List[Hook]] = None, + samplers: Optional[Union[Sampler, Dict[str, Sampler]]] = None, efficient_tuners: List[Dict] = None, **kwargs): @@ -132,6 +133,7 @@ class EpochBasedTrainer(BaseTrainer): self.train_dataloader = None self.eval_dataloader = None self.data_loader = None + self._samplers = samplers if isinstance(model, str): third_party = kwargs.get(ThirdParty.KEY) @@ -224,9 +226,6 @@ class EpochBasedTrainer(BaseTrainer): # Please check the DDPHook and MegatronHook for details. self.parallel_groups = {} - # Clear the Hook overload functions to avoid duplication. - Hook.clear_strategies() - if self.launcher is not None and not self.cfg.safe_get( 'train.hooks.DDPHook'): # A logic to fit the current code @@ -681,6 +680,7 @@ class EpochBasedTrainer(BaseTrainer): self.train_dataloader = self.get_train_dataloader() self.data_loader = self.train_dataloader self.register_optimizers_hook() + self.register_processors() self.print_hook_info() self.set_checkpoint_file_to_hook(checkpoint_path, load_all_state, kwargs.get('strict', False)) @@ -720,6 +720,7 @@ class EpochBasedTrainer(BaseTrainer): strict(`boolean`): If strict, any unmatched keys will cause an error. """ + self.register_processors() self.print_hook_info() if checkpoint_path is not None: from modelscope.trainers.hooks import LoadCheckpointHook @@ -758,6 +759,7 @@ class EpochBasedTrainer(BaseTrainer): kwargs: strict(`boolean`): If strict, any unmatched keys will cause an error. """ + self.register_processors() self.print_hook_info() if checkpoint_path is not None: from modelscope.trainers.hooks import LoadCheckpointHook @@ -897,11 +899,18 @@ class EpochBasedTrainer(BaseTrainer): """ if self.train_dataset is None: raise 'The train_dataset cannot be None.' + + sampler_cfg = {} + if self._samplers is not None: + sampler_cfg['sampler'] = self._samplers[ + ConfigKeys.train] if isinstance(self._samplers, + dict) else self._samplers data_loader = self._build_dataloader_with_dataset( self.train_dataset, dist=self._dist, seed=self._seed, collate_fn=self.train_data_collator, + **sampler_cfg, **self.cfg.train.get('dataloader', {})) return data_loader @@ -915,6 +924,11 @@ class EpochBasedTrainer(BaseTrainer): if self.eval_dataset is None: raise 'The eval_dataset cannot be None.' + sampler_cfg = {} + if self._samplers is not None: + sampler_cfg['sampler'] = self._samplers[ + ConfigKeys.val] if isinstance(self._samplers, + dict) else self._samplers default_config = {'shuffle': False} default_config.update(self.cfg.evaluation.get('dataloader', {})) data_loader = self._build_dataloader_with_dataset( @@ -922,6 +936,7 @@ class EpochBasedTrainer(BaseTrainer): dist=self._dist, seed=self._seed, collate_fn=self.eval_data_collator, + **sampler_cfg, **default_config) return data_loader @@ -938,6 +953,11 @@ class EpochBasedTrainer(BaseTrainer): mode=ModeKeys.EVAL, preprocessor=self.eval_preprocessor) + sampler_cfg = {} + if self._samplers is not None: + sampler_cfg['sampler'] = self._samplers[ + ConfigKeys.val] if isinstance(self._samplers, + dict) else self._samplers default_config = {'shuffle': False} default_config.update(self.cfg.evaluation.get('dataloader', {})) data_loader = self._build_dataloader_with_dataset( @@ -945,6 +965,7 @@ class EpochBasedTrainer(BaseTrainer): dist=self._dist, seed=self._seed, collate_fn=self.eval_data_collator, + **sampler_cfg, **default_config) return data_loader @@ -1132,13 +1153,19 @@ class EpochBasedTrainer(BaseTrainer): batch_size = batch_size_per_gpu num_workers = workers_per_gpu - if dist and not isinstance(dataset, torch.utils.data.IterableDataset): - sampler = DistributedSampler( - dataset, num_replicas=world_size, rank=rank, shuffle=shuffle) - else: - sampler = None - if not isinstance(dataset, torch.utils.data.IterableDataset): - kwargs['shuffle'] = shuffle + sampler = kwargs.pop('sampler', None) + if sampler is None: + if dist and not isinstance(dataset, + torch.utils.data.IterableDataset): + sampler = DistributedSampler( + dataset, + num_replicas=world_size, + rank=rank, + shuffle=shuffle) + else: + sampler = None + if not isinstance(dataset, torch.utils.data.IterableDataset): + kwargs['shuffle'] = shuffle batch_sampler = None @@ -1169,7 +1196,6 @@ class EpochBasedTrainer(BaseTrainer): """ Training loop used by `EpochBasedTrainer.train()` """ self.invoke_hook(TrainerStages.before_run) - kwargs = {} self.model.train() for _ in range(self._epoch, self._max_epochs): self.invoke_hook(TrainerStages.before_train_epoch) @@ -1181,7 +1207,7 @@ class EpochBasedTrainer(BaseTrainer): self.data_batch = data_batch self._inner_iter = i self.invoke_hook(TrainerStages.before_train_iter) - self.train_step(self.model, data_batch, **kwargs) + self.train_step(self.model, data_batch) self.invoke_hook(TrainerStages.after_train_iter) # Value changed after the hooks are invoked, do not move them above the invoke_hook code. del self.data_batch @@ -1320,12 +1346,17 @@ class EpochBasedTrainer(BaseTrainer): hooks = [] for cfg_i in hook_cfg: hook = build_from_cfg(cfg_i, HOOKS) - if hasattr(hook, 'register_strategy'): - hook.register_strategy() self.register_hook(hook) hooks.append(hook) return hooks + def register_processors(self): + """Register processors to hooks + """ + for hook in self.hooks: + if hasattr(hook, 'register_processor'): + hook.register_processor(self) + def get_hook(self, cls): return [h for h in self._hooks if h.__class__ == cls] @@ -1381,14 +1412,7 @@ class EpochBasedTrainer(BaseTrainer): info += '\n -------------------- ' stage_hook_infos.append(info) stage_hook_infos = '\n'.join(stage_hook_infos) - - strategy_info = '\n --- Hook strategies info --- \n' - for consumer, methods in Hook._strategies.items(): - strategy_info += f'Method: {consumer} ' \ - f'replaced by: ' \ - f'{[method.__self__.__class__.__name__ + "." + method.__name__ for method in methods]}\n' - strategy_info += '\n --- Hook strategies info end --- \n' - return stage_hook_infos + strategy_info + return stage_hook_infos def worker_init_fn(worker_id, num_workers, rank, seed): diff --git a/modelscope/trainers/training_args.py b/modelscope/trainers/training_args.py index a8eac217..260f2eb9 100644 --- a/modelscope/trainers/training_args.py +++ b/modelscope/trainers/training_args.py @@ -1,108 +1,544 @@ # Copyright (c) Alibaba, Inc. and its affiliates. - import re -from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentParser +from copy import deepcopy from dataclasses import dataclass, field, fields -from functools import partial -from typing import Any, Dict, List, Tuple, Union +from typing import List, Union -from modelscope.trainers.default_config import DEFAULT_CONFIG -from modelscope.utils.config import Config, ConfigDict -from modelscope.utils.hub import read_config +import addict +import json + +from modelscope.trainers.cli_argument_parser import CliArgumentParser +from modelscope.utils.config import Config -def get_flatten_value(config: Config, metadata: Dict, exclusions=None): - cfg_node = metadata['cfg_node'] - if exclusions is None: - exclusions = [] - - values = config.safe_get(cfg_node) - if isinstance(values, dict): - param_map = [] - for key, value in values.items(): - if key in exclusions or not isinstance(value, - (str, int, float, bool)): - continue - value = add_quotes_for_str(value) - param_map.append(f'{key}={value}') - return ','.join(param_map) - else: - return values - - -def set_flatten_value(config: Config, values: Union[str, List[str]], - metadata: Dict): - cfg_node = metadata['cfg_node'] - if values is None: - return config - +def set_flatten_value(values: Union[str, List[str]]): pairs = values.split(',') if isinstance(values, str) else values - for kv in pairs: + _params = {} + for kv in pairs or []: if len(kv.strip()) == 0: continue key, value = kv.split('=') - value = parse_value(value) - config.merge_from_dict({cfg_node + '.' + key: value}) - return config + _params[key] = parse_value(value) + return _params -def get_base_hook_args(config: Config, metadata: Dict): - cfg_node = metadata['cfg_node'] - hook_type = metadata['hook_type'] - key = metadata['key'] - value = config.safe_get(cfg_node) - if value is None: - return get_hook_param(config, hook_type, key) - else: - return True if key == 'type' else value +@dataclass +class DatasetArgs: + + train_dataset_name: str = field( + default=None, + metadata={ + 'help': + 'The dataset name used for training, can be an id in the datahub or a local dir', + }) + + val_dataset_name: str = field( + default=None, + metadata={ + 'help': + 'The subset name used for evaluating, can be an id in the datahub or a local dir', + }) + + train_subset_name: str = field( + default=None, + metadata={ + 'help': 'The subset name used for training, can be None', + }) + + val_subset_name: str = field( + default=None, + metadata={ + 'help': 'The subset name used for evaluating, can be None', + }) + + train_dataset_namespace: str = field( + default=None, + metadata={ + 'help': 'The dataset namespace used for training', + }) + + val_dataset_namespace: str = field( + default=None, + metadata={ + 'help': 'The dataset namespace used for evaluating', + }) + + dataset_json_file: str = field( + default=None, + metadata={ + 'help': + 'The json file to parse all datasets from, used in a complex dataset scenario,' + 'the json format should be like:' + ''' + [ + { + "dataset": { + # All args used in the MsDataset.load function + "dataset_name": "xxx", + ... + }, + # All columns used, mapping the column names in each dataset in same names. + "column_mapping": { + "text1": "sequence1", + "text2": "sequence2", + "label": "label", + }, + # float or str, float means to split the dataset into train/val, + # or just str(train/val) + "split": 0.8, + } + ] + ''', + }) -def set_base_hook_args(config: Config, value: Any, metadata: Dict): - cfg_node = metadata['cfg_node'] - hook_type = metadata['hook_type'] - key = metadata['key'] - if 'hooks' in config.train: - config.train.hooks = [ - hook for hook in config.train.hooks if hook['type'] != hook_type - ] - if key == 'type': - if value and config.safe_get(cfg_node) is None: - config.merge_from_dict({cfg_node: {}}) - else: - config.merge_from_dict({cfg_node: value}) +@dataclass +class ModelArgs: + task: str = field( + default=None, + metadata={ + 'help': 'The task code to be used', + 'cfg_node': 'task' + }) + + model: str = field( + default=None, metadata={ + 'help': 'A model id or model dir', + }) + + model_type: str = field( + default=None, + metadata={ + 'help': + 'The mode type, if load_model_config is False, user need to fill this field', + 'cfg_node': 'model.type' + }) -def get_strategy(config: Config, - metadata: Dict, - value_pair: Tuple[str] = ('by_epoch', 'by_step')): - flag = get_base_hook_args(config, metadata) - if flag is None: +@dataclass +class TrainArgs: + + seed: int = field( + default=42, metadata={ + 'help': 'The random seed', + }) + + per_device_train_batch_size: int = field( + default=16, + metadata={ + 'cfg_node': 'train.dataloader.batch_size_per_gpu', + 'help': + 'The `batch_size_per_gpu` argument for the train dataloader', + }) + + train_data_worker: int = field( + default=0, + metadata={ + 'cfg_node': 'train.dataloader.workers_per_gpu', + 'help': 'The `workers_per_gpu` argument for the train dataloader', + }) + + train_shuffle: bool = field( + default=False, + metadata={ + 'cfg_node': 'train.dataloader.shuffle', + 'help': 'The `shuffle` argument for the train dataloader', + }) + + train_drop_last: bool = field( + default=False, + metadata={ + 'cfg_node': 'train.dataloader.drop_last', + 'help': 'The `drop_last` argument for the train dataloader', + }) + + per_device_eval_batch_size: int = field( + default=16, + metadata={ + 'cfg_node': 'evaluation.dataloader.batch_size_per_gpu', + 'help': + 'The `batch_size_per_gpu` argument for the eval dataloader', + }) + + eval_data_worker: int = field( + default=0, + metadata={ + 'cfg_node': 'evaluation.dataloader.workers_per_gpu', + 'help': 'The `workers_per_gpu` argument for the eval dataloader', + }) + + eval_shuffle: bool = field( + default=False, + metadata={ + 'cfg_node': 'evaluation.dataloader.shuffle', + 'help': 'The `shuffle` argument for the eval dataloader', + }) + + eval_drop_last: bool = field( + default=False, + metadata={ + 'cfg_node': 'evaluation.dataloader.drop_last', + 'help': 'The `drop_last` argument for the eval dataloader', + }) + + max_epochs: int = field( + default=5, + metadata={ + 'cfg_node': 'train.max_epochs', + 'help': 'The training epochs', + }) + + work_dir: str = field( + default='./train_target', + metadata={ + 'cfg_node': 'train.work_dir', + 'help': 'The directory to save models and logs', + }) + + lr: float = field( + default=5e-5, + metadata={ + 'cfg_node': 'train.optimizer.lr', + 'help': 'The learning rate of the optimizer', + }) + + lr_scheduler: str = field( + default='LinearLR', + metadata={ + 'cfg_node': 'train.lr_scheduler.type', + 'help': 'The lr_scheduler type in torch', + }) + + optimizer: str = field( + default='AdamW', + metadata={ + 'cfg_node': 'train.optimizer.type', + 'help': 'The optimizer type in PyTorch, like `AdamW`', + }) + + optimizer_params: str = field( + default=None, + metadata={ + 'cfg_node': 'train.optimizer', + 'help': 'The optimizer params', + 'cfg_setter': set_flatten_value, + }) + + lr_scheduler_params: str = field( + default=None, + metadata={ + 'cfg_node': 'train.lr_scheduler', + 'help': 'The lr scheduler params', + 'cfg_setter': set_flatten_value, + }) + + lr_strategy: str = field( + default='by_epoch', + metadata={ + 'cfg_node': 'train.lr_scheduler.options.lr_strategy', + 'help': 'The lr decay strategy', + 'choices': ['by_epoch', 'by_step', 'no'], + }) + + local_rank: int = field( + default=0, metadata={ + 'help': 'The local rank', + }) + + logging_interval: int = field( + default=5, + metadata={ + 'help': 'The interval of iter of logging information', + 'cfg_node': 'train.logging.interval', + }) + + eval_strategy: str = field( + default='by_epoch', + metadata={ + 'help': 'Eval strategy, can be `by_epoch` or `by_step` or `no`', + 'cfg_node': 'evaluation.period.eval_strategy', + 'choices': ['by_epoch', 'by_step', 'no'], + }) + + eval_interval: int = field( + default=1, + metadata={ + 'help': 'Eval interval', + 'cfg_node': 'evaluation.period.interval', + }) + + eval_metrics: str = field( + default=None, + metadata={ + 'help': 'The metric name for evaluation', + 'cfg_node': 'evaluation.metrics' + }) + + save_strategy: str = field( + default='by_epoch', + metadata={ + 'help': + 'Checkpointing strategy, can be `by_epoch` or `by_step` or `no`', + 'cfg_node': 'train.checkpoint.period.save_strategy', + 'choices': ['by_epoch', 'by_step', 'no'], + }) + + save_interval: int = field( + default=1, + metadata={ + 'help': + 'The interval of epoch or iter of saving checkpoint period', + 'cfg_node': 'train.checkpoint.period.interval', + }) + + save_best_checkpoint: bool = field( + default=False, + metadata={ + 'help': + 'Save the checkpoint(if it\'s the best) after the evaluation.', + 'cfg_node': 'train.checkpoint.best.save_best', + }) + + metric_for_best_model: str = field( + default=None, + metadata={ + 'help': 'The metric used to measure the model.', + 'cfg_node': 'train.checkpoint.best.metric_key', + }) + + metric_rule_for_best_model: str = field( + default='max', + metadata={ + 'help': + 'The rule to measure the model with the metric, can be `max` or `min`', + 'cfg_node': 'train.checkpoint.best.rule', + }) + + max_checkpoint_num: int = field( + default=None, + metadata={ + 'help': + 'The max number of checkpoints to keep, older ones will be deleted.', + 'cfg_node': 'train.checkpoint.period.max_checkpoint_num', + }) + + max_checkpoint_num_best: int = field( + default=1, + metadata={ + 'help': + 'The max number of best checkpoints to keep, worse ones will be deleted.', + 'cfg_node': 'train.checkpoint.best.max_checkpoint_num', + }) + + push_to_hub: bool = field( + default=False, + metadata={ + 'help': 'Push to hub after each checkpointing', + 'cfg_node': 'train.checkpoint.period.push_to_hub', + }) + + repo_id: str = field( + default=None, + metadata={ + 'help': + 'The repo id in modelhub, usually the format is "group/model"', + 'cfg_node': 'train.checkpoint.period.hub_repo_id', + }) + + hub_token: str = field( + default=None, + metadata={ + 'help': + 'The modelhub token, you can also set the token to the env variable `MODELSCOPE_API_TOKEN`', + 'cfg_node': 'train.checkpoint.period.hub_token', + }) + + private_hub: bool = field( + default=True, + metadata={ + 'help': 'Upload to a private hub', + 'cfg_node': 'train.checkpoint.period.private_hub', + }) + + hub_revision: str = field( + default='master', + metadata={ + 'help': 'Which branch to commit to', + 'cfg_node': 'train.checkpoint.period.hub_revision', + }) + + push_to_hub_best: bool = field( + default=False, + metadata={ + 'help': 'Push to hub after each checkpointing', + 'cfg_node': 'train.checkpoint.best.push_to_hub', + }) + + repo_id_best: str = field( + default=None, + metadata={ + 'help': + 'The repo id in modelhub, usually the format is "group/model"', + 'cfg_node': 'train.checkpoint.best.hub_repo_id', + }) + + hub_token_best: str = field( + default=None, + metadata={ + 'help': + 'The modelhub token, you can also set the token to the env variable `MODELSCOPE_API_TOKEN`', + 'cfg_node': 'train.checkpoint.best.hub_token', + }) + + private_hub_best: bool = field( + default=True, + metadata={ + 'help': 'Upload to a private hub', + 'cfg_node': 'train.checkpoint.best.private_hub', + }) + + hub_revision_best: str = field( + default='master', + metadata={ + 'help': 'Which branch to commit to', + 'cfg_node': 'train.checkpoint.best.hub_revision', + }) + + +@dataclass(init=False) +class TrainingArgs(DatasetArgs, TrainArgs, ModelArgs): + + use_model_config: bool = field( + default=False, + metadata={ + 'help': + 'Use the configuration of the model, ' + 'default will only use the parameters in the CLI and the dataclass', + }) + + def __init__(self, **kwargs): + self.manual_args = list(kwargs.keys()) + for f in fields(self): + if f.name in kwargs: + setattr(self, f.name, kwargs[f.name]) + self._unknown_args = {} + + def parse_cli(self, parser_args=None): + """Construct a TrainingArg class by the parameters of CLI. + + Returns: + Self + """ + parser = CliArgumentParser(self) + args, unknown = parser.parse_known_args(parser_args) + unknown = [item for item in unknown if item not in ('\\', '\n')] + _unknown = {} + for i in range(0, len(unknown), 2): + _unknown[unknown[i].replace('-', '')] = parse_value(unknown[i + 1]) + args_dict = vars(args) + self.manual_args += parser.manual_args + + for key, value in deepcopy(args_dict).items(): + if key is not None and hasattr(self, key): + setattr(self, key, value) + return self + + def to_config(self, ignore_default_config=None): + """Convert the TrainingArgs to the `Config` + + Returns: + The Config, and extra parameters in dict. + """ + cfg = Config() + args_dict = addict.Dict() + + if ignore_default_config is None: + ignore_default_config = self.use_model_config + + for f in fields(self): + cfg_node = f.metadata.get('cfg_node') + cfg_setter = f.metadata.get('cfg_setter') or (lambda x: x) + if cfg_node is not None: + if f.name in self.manual_args or not ignore_default_config: + cfg.merge_from_dict( + {cfg_node: cfg_setter(getattr(self, f.name))}) + else: + args_dict[f.name] = getattr(self, f.name) + + cfg.merge_from_dict(self._unknown_args) + return cfg, args_dict + + def get_metadata(self, key): + _fields = fields(self) + for f in _fields: + if f.name == key: + return f return None - return value_pair[0] if flag else value_pair[1] -def set_strategy(config: Config, - value: Any, - metadata: Dict, - value_pair: Tuple[str] = ('by_epoch', 'by_step')): - set_base_hook_args(config, value == value_pair[0], metadata) +def build_dataset_from_file(filename): + """ + The filename format: + [ + { + "dataset": { + "dataset_name": "xxx", + ... + }, + "column_mapping": { + "text1": "sequence1", + "text2": "sequence2", + "label": "label", + } + "split": 0.8, + } + ] + """ + from modelscope import MsDataset + train_set = [] + eval_set = [] + with open(filename, 'r') as f: + ds_json = json.load(f) + for ds in ds_json: + dataset = MsDataset.load(**ds['dataset']).to_hf_dataset() + all_columns = dataset.column_names + keep_columns = ds['column_mapping'].keys() + remove_columns = [ + column for column in all_columns if column not in keep_columns + ] + from datasets import Features + from datasets import Value + from datasets import ClassLabel + features = [ + f for f in dataset.features.items() if f[0] in keep_columns + ] + new_features = {} + for f in features: + if isinstance(f[1], ClassLabel): + new_features[f[0]] = Value(f[1].dtype) + else: + new_features[f[0]] = f[1] + new_features = Features(new_features) + dataset = dataset.map( + lambda x: x, + remove_columns=remove_columns, + features=new_features).rename_columns(ds['column_mapping']) + split = ds['split'] + if isinstance(split, str): + assert split in ('train', 'val') + if split == 'train': + train_set.append(dataset) + else: + eval_set.append(dataset) + else: + assert isinstance(split, float) and 0 < split < 1 + ds_dict = dataset.train_test_split(train_size=split) + train_set.append(ds_dict['train']) + eval_set.append(ds_dict['test']) -def get_hook_param(config, hook_type: str, key='type'): - hooks = config.safe_get('train.hooks', []) - _hooks = list(filter(lambda hook: hook['type'] == hook_type, hooks)) - if key == 'type': - return len(_hooks) > 0 - elif len(_hooks) > 0: - return getattr(_hooks[0], key, None) - return None - - -def add_quotes_for_str(value: Union[str, float, bool, None]) -> str: - if isinstance(value, str): - return f'"{value}"' - else: - return str(value) + from datasets import concatenate_datasets + return concatenate_datasets(train_set), concatenate_datasets(eval_set) def parse_value(value: str) -> Union[str, float, bool, None]: @@ -126,714 +562,3 @@ def parse_value(value: str) -> Union[str, float, bool, None]: return float(value) else: return value - - -@dataclass -class TrainingArgs: - model: str = field( - default=None, metadata={ - 'help': 'A model id or model dir', - }) - - seed: int = field( - default=42, metadata={ - 'help': 'The random seed', - }) - - task: str = field( - default=None, - metadata={ - 'help': 'The task code to be used', - 'cfg_node': 'task' - }) - - dataset_name: str = field( - default=None, metadata={ - 'help': 'The dataset name', - }) - - subset_name: str = field( - default=None, metadata={ - 'help': 'The subset name of the dataset', - }) - - train_dataset_name: str = field( - default=None, metadata={ - 'help': 'The train dataset name', - }) - - val_dataset_name: str = field( - default=None, metadata={ - 'help': 'The validation dataset name', - }) - - per_device_train_batch_size: int = field( - default=None, - metadata={ - 'cfg_node': 'train.dataloader.batch_size_per_gpu', - 'help': 'The training batch size per GPU', - }) - - train_data_worker: int = field( - default=0, - metadata={ - 'cfg_node': 'train.dataloader.workers_per_gpu', - 'help': 'The number of data workers for train dataloader', - }) - - train_shuffle: bool = field( - default=None, - metadata={ - 'cfg_node': 'train.dataloader.shuffle', - 'help': 'Shuffle the train dataset or not', - }) - - train_drop_last: bool = field( - default=None, - metadata={ - 'cfg_node': - 'train.dataloader.drop_last', - 'help': - 'Whether to drop out the last set of data in the train_dataset', - }) - - per_device_eval_batch_size: int = field( - default=None, - metadata={ - 'cfg_node': 'evaluation.dataloader.batch_size_per_gpu', - 'help': 'The eval batch size per GPU', - }) - - eval_data_worker: int = field( - default=0, - metadata={ - 'cfg_node': 'evaluation.dataloader.workers_per_gpu', - 'help': 'The number of data workers for eval dataloader', - }) - - eval_shuffle: bool = field( - default=None, - metadata={ - 'cfg_node': 'evaluation.dataloader.shuffle', - 'help': 'Shuffle the eval dataset or not', - }) - - eval_drop_last: bool = field( - default=None, - metadata={ - 'cfg_node': 'evaluation.dataloader.drop_last', - 'help': - 'Whether to drop out the last set of data in the eval_dataset', - }) - - max_epochs: int = field( - default=None, - metadata={ - 'cfg_node': 'train.max_epochs', - 'help': 'The training epochs', - }) - - work_dir: str = field( - default=None, - metadata={ - 'cfg_node': 'train.work_dir', - 'help': 'The training dir to save models and logs', - }) - - lr: float = field( - default=None, - metadata={ - 'cfg_node': 'train.optimizer.lr', - 'help': 'The learning rate of the optimizer', - }) - - optimizer: str = field( - default=None, - metadata={ - 'cfg_node': 'train.optimizer.type', - 'help': 'The optimizer type', - }) - - optimizer_params: str = field( - default=None, - metadata={ - 'cfg_node': - 'train.optimizer', - 'cfg_getter': - partial(get_flatten_value, exclusions=['type', 'lr', 'options']), - 'cfg_setter': - set_flatten_value, - 'help': - 'The optimizer init params except `lr`', - }) - - lr_scheduler_params: str = field( - default=None, - metadata={ - 'cfg_node': - 'train.lr_scheduler', - 'cfg_getter': - partial(get_flatten_value, exclusions=['type', 'lr', 'options']), - 'cfg_setter': - set_flatten_value, - 'help': - 'The lr_scheduler init params', - }) - - local_rank: int = field( - default=0, metadata={ - 'help': 'The training local rank', - }) - - save_ckpt: bool = field( - default=True, - metadata={ - 'help': - 'Periodically save checkpoint when True, corresponding to CheckpointHook', - 'cfg_node': 'train.checkpoint.period', - 'hook_type': 'CheckpointHook', - 'key': 'type', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - save_ckpt_best: bool = field( - default=None, - metadata={ - 'help': - 'Save best checkpoint when True, corresponding to BestCkptSaverHook', - 'cfg_node': 'train.checkpoint.best', - 'hook_type': 'BestCkptSaverHook', - 'key': 'type', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - evaluate: bool = field( - default=True, - metadata={ - 'help': 'Evaluate when True, corresponding to EvaluationHook', - 'cfg_node': 'evaluation.period', - 'hook_type': 'EvaluationHook', - 'key': 'type', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - save_ckpt_strategy: str = field( - default=None, - metadata={ - 'help': 'Periodically save checkpoint by epoch or by step' - 'use with `CheckpointHook`, can be `by_epoch` or `by_step`', - 'cfg_node': 'train.checkpoint.period.by_epoch', - 'hook_type': 'CheckpointHook', - 'key': 'by_epoch', - 'choices': ['by_epoch', 'by_step'], - 'cfg_getter': get_strategy, - 'cfg_setter': set_strategy, - }) - - save_ckpt_best_strategy: str = field( - default=None, - metadata={ - 'help': 'Save best checkpoint by epoch or by step' - 'use with `BestCkptSaverHook`, can be `by_epoch` or `by_step`', - 'cfg_node': 'train.checkpoint.best.by_epoch', - 'hook_type': 'BestCkptSaverHook', - 'key': 'by_epoch', - 'choices': ['by_epoch', 'by_step'], - 'cfg_getter': get_strategy, - 'cfg_setter': set_strategy, - }) - - push_to_hub: bool = field( - default=None, - metadata={ - 'help': - 'Push to hub after one checkpoint saved by CheckpointHook in the local disk', - 'cfg_node': 'train.checkpoint.period.push_to_hub', - 'hook_type': 'CheckpointHook', - 'key': 'push_to_hub', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - model_id_with_org: str = field( - default=None, - metadata={ - 'help': - 'The repo id in modelhub, usually it\'s like "group/model"', - 'cfg_node': 'train.checkpoint.period.model_id_with_org', - 'hook_type': 'CheckpointHook', - 'key': 'model_id_with_org', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - hub_token: str = field( - default=None, - metadata={ - 'help': - 'The token to push to hub, you can also set the token to the env variable `MODELSCOPE_API_TOKEN`', - 'cfg_node': 'train.checkpoint.period.hub_token', - 'hook_type': 'CheckpointHook', - 'key': 'hub_token', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - private_hub: bool = field( - default=None, - metadata={ - 'help': 'Upload to a private hub', - 'cfg_node': 'train.checkpoint.period.private_hub', - 'hook_type': 'CheckpointHook', - 'key': 'private_hub', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - push_to_hub_best_model: bool = field( - default=None, - metadata={ - 'help': - 'Push to hub after one checkpoint saved by BestCkptSaverHook in the local disk', - 'cfg_node': 'train.checkpoint.best.push_to_hub', - 'hook_type': 'BestCkptSaverHook', - 'key': 'push_to_hub', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - model_id_with_org_best_model: str = field( - default=None, - metadata={ - 'help': - 'The repo id in modelhub, usually it\'s like "group/model"', - 'cfg_node': 'train.checkpoint.best.model_id_with_org', - 'hook_type': 'BestCkptSaverHook', - 'key': 'model_id_with_org', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - hub_token_best_model: str = field( - default=None, - metadata={ - 'help': - 'The token to push to hub, you can also set the token to the env variable `MODELSCOPE_API_TOKEN`', - 'cfg_node': 'train.checkpoint.best.hub_token', - 'hook_type': 'BestCkptSaverHook', - 'key': 'hub_token', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - private_hub_best_model: bool = field( - default=None, - metadata={ - 'help': 'Upload to a private hub', - 'cfg_node': 'train.checkpoint.best.private_hub', - 'hook_type': 'BestCkptSaverHook', - 'key': 'private_hub', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - ckpt_period_interval: int = field( - default=1, - metadata={ - 'help': - 'The interval of epoch or iter of saving checkpoint period', - 'cfg_node': 'train.checkpoint.period.interval', - 'hook_type': 'CheckpointHook', - 'key': 'interval', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - ckpt_best_interval: int = field( - default=None, - metadata={ - 'help': 'The interval of epoch or iter of saving checkpoint best', - 'cfg_node': 'train.checkpoint.best.interval', - 'hook_type': 'BestCkptSaverHook', - 'key': 'interval', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - metric_for_best_model: str = field( - default=None, - metadata={ - 'help': - 'Which metric key to judge the checkpoint is better or not, use with `BestCkptSaverHook`, ' - 'please make sure this key is returned by the `evaluation_metrics` classes', - 'cfg_node': - 'train.checkpoint.best.metric_key', - 'hook_type': - 'BestCkptSaverHook', - 'key': - 'metric_key', - 'cfg_getter': - get_base_hook_args, - 'cfg_setter': - set_base_hook_args, - }) - - metric_rule_for_best_model: str = field( - default=None, - metadata={ - 'help': - 'Which rule to compare the value of `checkpoint_saving_metric`, ' - 'use with `BestCkptSaverHook`, can be `max` or `min`', - 'cfg_node': - 'train.checkpoint.best.rule', - 'hook_type': - 'BestCkptSaverHook', - 'key': - 'rule', - 'cfg_getter': - get_base_hook_args, - 'cfg_setter': - set_base_hook_args, - }) - - save_ckpt_peroid_limit: int = field( - default=None, - metadata={ - 'help': - 'The max saving number of checkpoint, older checkpoints will be deleted.', - 'cfg_node': 'train.checkpoint.period.max_checkpoint_num', - 'hook_type': 'CheckpointHook', - 'key': 'max_checkpoint_num', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - save_ckpt_best_limit: int = field( - default=None, - metadata={ - 'help': - 'The max saving number of checkpoint, worse checkpoints will be deleted.', - 'cfg_node': 'train.checkpoint.best.max_checkpoint_num', - 'hook_type': 'BestCkptSaverHook', - 'key': 'max_checkpoint_num', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - logging_interval: int = field( - default=None, - metadata={ - 'help': 'The interval of iter of logging information', - 'cfg_node': 'train.logging.interval', - 'hook_type': 'TextLoggerHook', - 'key': 'interval', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - eval_strategy: str = field( - default=None, - metadata={ - 'help': 'Evaluate model by epoch or by step' - 'use with `EvaluationHook`, can be `by_epoch` or `by_step`', - 'cfg_node': 'evaluation.period.by_epoch', - 'hook_type': 'EvaluationHook', - 'key': 'by_epoch', - 'choices': ['by_epoch', 'by_step'], - 'cfg_getter': get_strategy, - 'cfg_setter': set_strategy, - }) - - eval_interval: int = field( - default=None, - metadata={ - 'help': 'Evaluation interval by epoch or iter', - 'cfg_node': 'evaluation.period.interval', - 'hook_type': 'EvaluationHook', - 'key': 'interval', - 'cfg_getter': get_base_hook_args, - 'cfg_setter': set_base_hook_args, - }) - - eval_metrics: str = field( - default=None, - metadata={ - 'help': 'The metric module name used in evaluation', - 'cfg_node': 'evaluation.metrics' - }) - - @classmethod - def from_cli(cls, parser_args=None, **extra_kwargs): - """Construct a TrainingArg class by the parameters of CLI. - - Args: - **extra_kwargs: Extra args which can be defined in code. - - Returns: - The output TrainingArg class with the parameters from CLI. - """ - self = cls(**extra_kwargs) - parser = CliArgumentParser(self) - args, unknown = parser.parse_known_args(parser_args) - unknown = [item for item in unknown if item not in ('\\', '\n')] - _unknown = {} - for i in range(0, len(unknown), 2): - _unknown[unknown[i].replace('-', '')] = parse_value(unknown[i + 1]) - cfg_dict = vars(args) - - if args.model is not None: - try: - cfg = read_config(args.model) - except Exception as e: - print('Read config failed with error:', e) - else: - self = cls.from_config(cfg, **extra_kwargs) - for key, value in cfg_dict.items(): - if key is not None and hasattr(self, - key) and key in parser.manual_args: - setattr(self, key, value) - self.extra_args = _unknown - return self - - def to_args(self): - """Convert the TrainingArg class to key-value pairs. - - Returns: The key-value pair. - - """ - _args = {} - for f in fields(self): - _args[f.name] = getattr(self, f.name) - return _args - - @classmethod - def from_config(cls, config=DEFAULT_CONFIG, **kwargs): - """Construct the TrainingArg class by a `Config` class. - - Args: - config: The Config class. By default, `DEFAULT_CONFIG` is used. - **kwargs: Extra args which can be defined in code. - - Returns: The output TrainingArg class with the parameters from the config. - - """ - - self = cls(**kwargs) - for f in fields(self): - if 'cfg_node' in f.metadata and getattr(self, f.name) is None: - self._to_field(f, config) - return self - - def _to_field(self, f, config): - assert 'cfg_node' in f.metadata - if 'cfg_getter' in f.metadata: - cfg_getter = f.metadata['cfg_getter'] - setattr(self, f.name, cfg_getter(config, f.metadata)) - else: - cfg_node = f.metadata['cfg_node'] - setattr(self, f.name, config.safe_get(cfg_node)) - - def _to_config(self, f, config: Config): - assert 'cfg_node' in f.metadata - value = getattr(self, f.name) - if 'cfg_setter' in f.metadata: - cfg_setter = f.metadata['cfg_setter'] - config = cfg_setter(config, value, f.metadata) - else: - cfg_node = f.metadata['cfg_node'] - if isinstance(cfg_node, str): - cfg_node = [cfg_node] - for _node in cfg_node: - config.merge_from_dict({_node: value}) - return config - - def __call__(self, cfg: Config): - for f in fields(self): - if 'cfg_node' not in f.metadata: - continue - - value = getattr(self, f.name) - if value is not None: - self._to_config(f, cfg) - if hasattr(self, 'extra_args'): - cfg.merge_from_dict(self.extra_args) - else: - self._to_field(f, cfg) - return cfg - - -class CliArgumentParser(ArgumentParser): - """ Argument Parser to define and parse command-line args for training. - - Args: - training_args (TrainingArgs): dict or list of dict which defines different - paramters for training. - """ - - def __init__(self, training_args: TrainingArgs = None, **kwargs): - if 'formatter_class' not in kwargs: - kwargs['formatter_class'] = ArgumentDefaultsHelpFormatter - super().__init__(**kwargs) - self.training_args = training_args - self.define_args() - - def get_manual_args(self, args): - return [arg[2:] for arg in args if arg.startswith('--')] - - def _parse_known_args(self, args: List = None, namespace=None): - self.model_id = namespace.model if namespace is not None else None - if '--model' in args: - self.model_id = args[args.index('--model') + 1] - self.manual_args = self.get_manual_args(args) - return super()._parse_known_args(args, namespace) - - def print_help(self, file=None): - config = DEFAULT_CONFIG - if self.model_id is not None: - try: - config = read_config(self.model_id) - except Exception as e: - print('Read config failed with error:', e) - - if config is not None: - for action_group in self._optionals._group_actions: - if hasattr(self.training_args, action_group.dest): - value = getattr(self.training_args, action_group.dest) - f = {f.name: f - for f in fields(self.training_args) - }.get(action_group.dest) - if value is not None: - action_group.default = value - elif 'cfg_node' in f.metadata: - cfg_node = f.metadata['cfg_node'] - if isinstance(cfg_node, str): - cfg_node = [cfg_node] - - assert isinstance(cfg_node, (list, tuple)) - if isinstance(cfg_node[0], str): - action_group.default = config.safe_get(cfg_node[0]) - else: - action_group.default = cfg_node[0](config) - return super().print_help(file) - - def define_args(self): - if self.training_args is not None: - for f in fields(self.training_args): - arg_name = f.name - arg_attr = getattr(self.training_args, f.name) - name = f'--{arg_name}' - kwargs = dict(type=f.type, help=f.metadata['help']) - kwargs['default'] = arg_attr - - if 'choices' in f.metadata: - kwargs['choices'] = f.metadata['choices'] - - kwargs['action'] = SingleAction - self.add_argument(name, **kwargs) - - -class DictAction(Action): - """ - argparse action to split an argument into KEY=VALUE form - on the first = and append to a dictionary. List options can - be passed as comma separated values, i.e 'KEY=V1,V2,V3', or with explicit - brackets, i.e. 'KEY=[V1,V2,V3]'. It also support nested brackets to build - list/tuple values. e.g. 'KEY=[(V1,V2),(V3,V4)]' - """ - - @staticmethod - def parse_int_float_bool_str(val): - try: - return int(val) - except ValueError: - pass - try: - return float(val) - except ValueError: - pass - if val.lower() in ['true', 'false']: - return val.lower() == 'true' - if val == 'None': - return None - return val - - @staticmethod - def parse_iterable(val): - """Parse iterable values in the string. - All elements inside '()' or '[]' are treated as iterable values. - Args: - val (str): Value string. - Returns: - list | tuple: The expanded list or tuple from the string. - Examples: - >>> DictAction._parse_iterable('1,2,3') - [1, 2, 3] - >>> DictAction._parse_iterable('[a, b, c]') - ['a', 'b', 'c'] - >>> DictAction._parse_iterable('[(1, 2, 3), [a, b], c]') - [(1, 2, 3), ['a', 'b'], 'c'] - """ - - def find_next_comma(string): - """Find the position of next comma in the string. - If no ',' is found in the string, return the string length. All - chars inside '()' and '[]' are treated as one element and thus ',' - inside these brackets are ignored. - """ - assert (string.count('(') == string.count(')')) and ( - string.count('[') - == string.count(']')), f'Imbalanced brackets exist in {string}' - end = len(string) - for idx, char in enumerate(string): - pre = string[:idx] - # The string before this ',' is balanced - if ((char == ',') and (pre.count('(') == pre.count(')')) - and (pre.count('[') == pre.count(']'))): - end = idx - break - return end - - # Strip ' and " characters and replace whitespace. - val = val.strip('\'\"').replace(' ', '') - is_tuple = False - if val.startswith('(') and val.endswith(')'): - is_tuple = True - val = val[1:-1] - elif val.startswith('[') and val.endswith(']'): - val = val[1:-1] - elif ',' not in val: - # val is a single value - return DictAction.parse_int_float_bool_str(val) - - values = [] - while len(val) > 0: - comma_idx = find_next_comma(val) - element = DictAction.parse_iterable(val[:comma_idx]) - values.append(element) - val = val[comma_idx + 1:] - if is_tuple: - values = tuple(values) - return values - - def __call__(self, parser, namespace, values, option_string): - options = {} - for kv in values: - key, val = kv.split('=', maxsplit=1) - options[key] = self.parse_iterable(val) - setattr(namespace, self.dest, options) - - -class SingleAction(DictAction): - """ Argparse action to convert value to tuple or list or nested structure of - list and tuple, i.e 'V1,V2,V3', or with explicit brackets, i.e. '[V1,V2,V3]'. - It also support nested brackets to build list/tuple values. e.g. '[(V1,V2),(V3,V4)]' - """ - - def __call__(self, parser, namespace, value, option_string): - if isinstance(value, str): - setattr(namespace, self.dest, self.parse_iterable(value)) - else: - setattr(namespace, self.dest, value) diff --git a/modelscope/utils/ast_utils.py b/modelscope/utils/ast_utils.py index 374ada20..5cee374d 100644 --- a/modelscope/utils/ast_utils.py +++ b/modelscope/utils/ast_utils.py @@ -1,7 +1,6 @@ # Copyright (c) Alibaba, Inc. and its affiliates. import ast -import contextlib import hashlib import os import os.path as osp @@ -9,12 +8,11 @@ import time import traceback from functools import reduce from pathlib import Path -from typing import Generator, Union +from typing import Union import gast import json -from modelscope import __version__ from modelscope.fileio.file import LocalStorage from modelscope.metainfo import (CustomDatasets, Heads, Hooks, LR_Schedulers, Metrics, Models, Optimizers, Pipelines, @@ -574,6 +572,7 @@ file_scanner = FilesAstScanning() def _save_index(index, file_path, file_list=None, with_template=False): # convert tuple key to str key index[INDEX_KEY] = {str(k): v for k, v in index[INDEX_KEY].items()} + from modelscope.version import __version__ index[VERSION_KEY] = __version__ index[MD5_KEY], index[FILES_MTIME_KEY] = file_scanner.files_mtime_md5( file_list=file_list) @@ -682,6 +681,7 @@ def load_index( if not force_rebuild and os.path.exists(file_path): wrapped_index = _load_index(file_path) md5, files_mtime = file_scanner.files_mtime_md5(file_list=file_list) + from modelscope.version import __version__ if (wrapped_index[VERSION_KEY] == __version__): index = wrapped_index if (wrapped_index[MD5_KEY] != md5): diff --git a/modelscope/utils/checkpoint.py b/modelscope/utils/checkpoint.py index 64681db4..bbde6034 100644 --- a/modelscope/utils/checkpoint.py +++ b/modelscope/utils/checkpoint.py @@ -5,7 +5,6 @@ import os import re import time from collections import OrderedDict -from functools import partial from shutil import copytree, ignore_patterns, rmtree from typing import Callable, Dict, Optional, Union @@ -15,7 +14,6 @@ from torch import nn from torch.optim import Optimizer from torch.optim.lr_scheduler import _LRScheduler -from modelscope import __version__ from modelscope.fileio import File, LocalStorage from modelscope.utils.config import Config, JSONIteratorEncoder from modelscope.utils.constant import ConfigFields, ModelFile @@ -76,6 +74,7 @@ def save_checkpoint(model: torch.nn.Module, elif not isinstance(meta, dict): raise TypeError( f'meta must be a dict or None, but got {type(meta)}') + from modelscope import __version__ meta.update(modelscope=__version__, time=time.asctime()) if isinstance(model, torch.nn.parallel.DistributedDataParallel): diff --git a/tests/hub/test_hub_upload.py b/tests/hub/test_hub_upload.py index a8b90288..2a66cb8b 100644 --- a/tests/hub/test_hub_upload.py +++ b/tests/hub/test_hub_upload.py @@ -37,6 +37,7 @@ class HubUploadTest(unittest.TestCase): os.mkdir(self.finetune_path) os.system("echo '{}'>%s" % os.path.join(self.finetune_path, ModelFile.CONFIGURATION)) + os.environ['MODELSCOPE_TRAIN_ID'] = 'test-id' def tearDown(self): logger.info('TearDown') diff --git a/tests/trainers/cli/__init__.py b/tests/trainers/cli/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/trainers/cli/test_cli.py b/tests/trainers/cli/test_cli.py new file mode 100644 index 00000000..b9fb7539 --- /dev/null +++ b/tests/trainers/cli/test_cli.py @@ -0,0 +1,52 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import unittest + +import json + +from modelscope import MsDataset, TrainingArgs, build_dataset_from_file +from modelscope.utils.test_utils import test_level + + +class TestCli(unittest.TestCase): + + def setUp(self) -> None: + content = [{ + 'dataset': { + 'dataset_name': 'clue', + 'subset_name': 'cmnli', + 'split': 'train', + }, + 'column_mapping': { + 'sentence1': 'sentence1', + 'sentence2': 'sentence2', + 'label': 'label', + }, + 'split': 0.8, + }, { + 'dataset': { + 'dataset_name': 'glue', + 'subset_name': 'mnli', + 'split': 'validation_matched', + }, + 'column_mapping': { + 'premise': 'sentence1', + 'hypothesis': 'sentence2', + 'label': 'label', + }, + 'split': 'val', + }] + with open('./dataset.json', 'w') as f: + json.dump(content, f) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_merge_dataset_from_file(self): + dataset = MsDataset.load('clue', subset_name='cmnli', split='train') + dataset2 = MsDataset.load( + 'glue', subset_name='mnli', split='validation_matched') + training_args = TrainingArgs(dataset_json_file='./dataset.json') + train, test = build_dataset_from_file(training_args.dataset_json_file) + self.assertEqual(len(train) + len(test), len(dataset) + len(dataset2)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/trainers/hooks/test_lr_scheduler_hook.py b/tests/trainers/hooks/test_lr_scheduler_hook.py index cd28b055..432fb39a 100644 --- a/tests/trainers/hooks/test_lr_scheduler_hook.py +++ b/tests/trainers/hooks/test_lr_scheduler_hook.py @@ -105,6 +105,7 @@ class LrSchedulerHookTest(unittest.TestCase): train_dataloader = trainer._build_dataloader_with_dataset( trainer.train_dataset, **trainer.cfg.train.get('dataloader', {})) trainer.register_optimizers_hook() + trainer.register_processors() trainer._hooks = [ hook for hook in trainer._hooks if hook.__class__.__name__ not in ['CheckpointHook', 'TextLoggerHook', 'IterTimerHook'] @@ -177,6 +178,7 @@ class LrSchedulerHookTest(unittest.TestCase): train_dataloader = trainer._build_dataloader_with_dataset( trainer.train_dataset, **trainer.cfg.train.get('dataloader', {})) trainer.register_optimizers_hook() + trainer.register_processors() trainer._hooks = [ hook for hook in trainer._hooks if hook.__class__.__name__ not in ['CheckpointHook', 'TextLoggerHook', 'IterTimerHook'] @@ -365,6 +367,7 @@ class PlateauLrSchedulerHookTest(unittest.TestCase): trainer.train_dataloader = train_dataloader trainer.data_loader = train_dataloader trainer.register_optimizers_hook() + trainer.register_processors() trainer._hooks = [ hook for hook in trainer._hooks if hook.__class__.__name__ not in ['CheckpointHook', 'TextLoggerHook', 'IterTimerHook'] diff --git a/tests/trainers/hooks/test_optimizer_hook.py b/tests/trainers/hooks/test_optimizer_hook.py index b9899c36..ed0e202a 100644 --- a/tests/trainers/hooks/test_optimizer_hook.py +++ b/tests/trainers/hooks/test_optimizer_hook.py @@ -150,6 +150,7 @@ class TorchAMPOptimizerHookTest(unittest.TestCase): train_dataloader = trainer._build_dataloader_with_dataset( trainer.train_dataset, **trainer.cfg.train.get('dataloader', {})) trainer.register_optimizers_hook() + trainer.register_processors() trainer._hooks = [ hook for hook in trainer._hooks if hook.__class__.__name__ not in ['CheckpointHook', 'TextLoggerHook', 'IterTimerHook'] diff --git a/tests/trainers/test_trainer_with_nlp.py b/tests/trainers/test_trainer_with_nlp.py index ceb04e15..a736d4fa 100644 --- a/tests/trainers/test_trainer_with_nlp.py +++ b/tests/trainers/test_trainer_with_nlp.py @@ -9,6 +9,7 @@ import unittest import numpy as np import torch from packaging import version +from torch.utils.data import RandomSampler from modelscope.hub.snapshot_download import snapshot_download from modelscope.metainfo import Metrics @@ -204,12 +205,20 @@ class TestTrainerWithNlp(unittest.TestCase): cfg.preprocessor.val['label2id'] = {'0': 0, '1': 1} cfg.train.dataloader.batch_size_per_gpu = 2 cfg.train.hooks = [{ - 'type': 'BestCkptSaverHook', - 'interval': 1, - 'by_epoch': False, - 'metric_key': 'accuracy', - 'max_checkpoint_num': 4, - 'restore_best': True, + 'type': + 'BestCkptSaverHook', + 'interval': + 1, + 'by_epoch': + False, + 'output_dir': + os.path.join(self.tmp_dir, 'output_test_best'), + 'metric_key': + 'accuracy', + 'max_checkpoint_num': + 4, + 'restore_best': + True, }, { 'type': 'TextLoggerHook', 'interval': 1 @@ -270,7 +279,7 @@ class TestTrainerWithNlp(unittest.TestCase): os.path.join(self.tmp_dir, 'output', 'pytorch_model.bin'))) self.assertTrue( os.path.isfile( - os.path.join(self.tmp_dir, 'output_best', + os.path.join(self.tmp_dir, 'output_test_best', 'pytorch_model.bin'))) md51 = hashlib.md5( pathlib.Path( @@ -282,7 +291,7 @@ class TestTrainerWithNlp(unittest.TestCase): self.assertEqual(md51, md52) md51 = hashlib.md5( pathlib.Path( - os.path.join(self.tmp_dir, 'output_best', + os.path.join(self.tmp_dir, 'output_test_best', 'pytorch_model.bin')).read_bytes()).hexdigest() md52 = hashlib.md5( pathlib.Path( @@ -472,6 +481,34 @@ class TestTrainerWithNlp(unittest.TestCase): cache_path + '/pytorch_model.bin', saving_fn=saving_fn)) self.assertTrue(os.path.isfile(f'{tmp_dir}/predicts.txt')) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_trainer_with_custom_sampler(self): + tmp_dir = tempfile.TemporaryDirectory().name + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + model_id = 'damo/nlp_structbert_sentence-similarity_chinese-tiny' + cache_path = snapshot_download(model_id) + model = SbertForSequenceClassification.from_pretrained(cache_path) + + class CustomSampler(RandomSampler): + + pass + + kwargs = dict( + cfg_file=os.path.join(cache_path, ModelFile.CONFIGURATION), + model=model, + train_dataset=self.dataset, + eval_dataset=self.dataset, + samplers=CustomSampler(self.dataset), + work_dir=self.tmp_dir) + + trainer = build_trainer(default_args=kwargs) + trainer.train() + self.assertTrue( + type(trainer.train_dataloader.sampler) == CustomSampler) + self.assertTrue(type(trainer.eval_dataloader.sampler) == CustomSampler) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_trainer_with_prediction(self): tmp_dir = tempfile.TemporaryDirectory().name diff --git a/tests/trainers/test_training_args.py b/tests/trainers/test_training_args.py index 6e4d306e..e8f6d8a2 100644 --- a/tests/trainers/test_training_args.py +++ b/tests/trainers/test_training_args.py @@ -1,8 +1,8 @@ # Copyright (c) Alibaba, Inc. and its affiliates. import unittest -from modelscope.trainers.default_config import DEFAULT_CONFIG -from modelscope.trainers.training_args import CliArgumentParser, TrainingArgs +from modelscope import TrainingArgs +from modelscope.trainers.cli_argument_parser import CliArgumentParser from modelscope.utils.test_utils import test_level @@ -29,14 +29,14 @@ class TrainingArgsTest(unittest.TestCase): @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_flatten_args(self): - cfg = DEFAULT_CONFIG + training_args = TrainingArgs() input_args = [ '--optimizer_params', 'weight_decay=0.8,eps=1e-6,correct_bias=False', '--lr_scheduler_params', 'initial_lr=3e-5,niter_decay=1' ] - training_args = TrainingArgs.from_cli(input_args) - cfg = training_args(cfg) + training_args = training_args.parse_cli(input_args) + cfg, _ = training_args.to_config() self.assertAlmostEqual(cfg.train.optimizer.weight_decay, 0.8) self.assertAlmostEqual(cfg.train.optimizer.eps, 1e-6) self.assertFalse(cfg.train.optimizer.correct_bias)