Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2026-09-01 19:49:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

merge master

Browse Source
This commit is contained in:
xingjun.wang
2023-05-14 02:09:02 +08:00
parent bf2cd23aa2 febc0365de
commit 98e899b8c2
46 changed files with 2394 additions and 2104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/pytorch/image_classification/finetune_image_classification.py
View File

@@ -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

4
examples/pytorch/multi_modal_embedding/finetune_multi_modal_embedding.py
View File

@@ -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

57
examples/pytorch/text_classification/finetune_text_classification.py
View File

@@ -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)

5
examples/pytorch/text_classification/run_train.sh
View File

@@ -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 \

48
examples/pytorch/text_generation/finetune_text_generation.py
View File

@@ -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)

5
examples/pytorch/text_generation/run_train_gpt3.sh
View File

@@ -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

4
examples/pytorch/token_classification/finetune_token_classification.py
View File

@@ -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

79
modelscope/__init__.py
View File

@@ -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={},
)

19
modelscope/hub/api.py
View File

@@ -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(),

1
modelscope/hub/file_download.py
View File

@@ -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,

44
modelscope/hub/push_to_hub.py
View File

@@ -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)

4
modelscope/models/audio/tts/voice.py
View File

@@ -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:

2
modelscope/pipelines/__init__.py
View File

@@ -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

2
modelscope/preprocessors/__init__.py
View File

@@ -35,7 +35,7 @@ if TYPE_CHECKING:
TextErrorCorrectionPreprocessor, TextGenerationT5Preprocessor,
WordAlignmentPreprocessor, TextGenerationTransformersPreprocessor,
Tokenize, WordSegmentationBlankSetToLabelPreprocessor,
CodeGeeXPreprocessor, MGLMSummarizationPreprocessor,
MGLMSummarizationPreprocessor,
ZeroShotClassificationTransformersPreprocessor,
TextGenerationJiebaPreprocessor, SentencePiecePreprocessor,
DialogIntentPredictionPreprocessor, DialogModelingPreprocessor,

6
modelscope/trainers/__init__.py
View File

@@ -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

151
modelscope/trainers/cli_argument_parser.py Normal file
View File

@@ -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)

37
modelscope/trainers/default_config.py
View File

@@ -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.'

6
modelscope/trainers/hooks/__init__.py
View File

@@ -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']
}

2
modelscope/trainers/hooks/checkpoint/__init__.py Normal file
View File

@@ -0,0 +1,2 @@
from .checkpoint_hook import BestCkptSaverHook, CheckpointHook
from .load_checkpoint_hook import LoadCheckpointHook

435
modelscope/trainers/hooks/checkpoint/checkpoint_hook.py Normal file
View File

@@ -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)

276
modelscope/trainers/hooks/checkpoint/checkpoint_processor.py Normal file
View File

@@ -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:]

138
modelscope/trainers/hooks/checkpoint/load_checkpoint_hook.py Normal file
View File

@@ -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

749
modelscope/trainers/hooks/checkpoint_hook.py
View File

@@ -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:]

1
modelscope/trainers/hooks/compression/sparsity_hook.py
View File

@@ -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

3
modelscope/trainers/hooks/distributed/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
from .ddp_hook import DDPHook
from .deepspeed_hook import DeepspeedHook
from .megatron_hook import MegatronHook

6
modelscope/trainers/hooks/ddp_hook.py → modelscope/trainers/hooks/distributed/ddp_hook.py
View File

@@ -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)

145
modelscope/trainers/hooks/deepspeed_hook.py → modelscope/trainers/hooks/distributed/deepspeed_hook.py
View File

@@ -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

239
modelscope/trainers/hooks/megatron_hook.py → modelscope/trainers/hooks/distributed/megatron_hook.py
View File

@@ -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)

30
modelscope/trainers/hooks/early_stop_hook.py
View File

@@ -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):

30
modelscope/trainers/hooks/evaluation_hook.py
View File

@@ -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:

72
modelscope/trainers/hooks/hook.py
View File

@@ -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

132
modelscope/trainers/hooks/lr_scheduler_hook.py
View File

@@ -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):

72
modelscope/trainers/hooks/optimizer/apex_optimizer_hook.py
View File

@@ -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))

89
modelscope/trainers/hooks/optimizer/base.py
View File

@@ -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)

84
modelscope/trainers/hooks/optimizer/torch_optimizer_hook.py
View File

@@ -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))

72
modelscope/trainers/trainer.py
View File

@@ -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):

1318
modelscope/trainers/training_args.py
View File

File diff suppressed because it is too large Load Diff

6
modelscope/utils/ast_utils.py
View File

@@ -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):

3
modelscope/utils/checkpoint.py
View File

@@ -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):

1
tests/hub/test_hub_upload.py
View File

@@ -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')

0
tests/trainers/cli/__init__.py Normal file
View File

52
tests/trainers/cli/test_cli.py Normal file
View File

@@ -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()

3
tests/trainers/hooks/test_lr_scheduler_hook.py
View File

@@ -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']

1
tests/trainers/hooks/test_optimizer_hook.py
View File

@@ -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']

53
tests/trainers/test_trainer_with_nlp.py
View File

@@ -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

10
tests/trainers/test_training_args.py
View File

@@ -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)