mirror of
https://github.com/modelscope/modelscope.git
synced 2025-12-17 00:37:43 +01:00
91 lines
2.4 KiB
Python
91 lines
2.4 KiB
Python
import os
|
|
from dataclasses import dataclass, field
|
|
|
|
from modelscope import MsDataset, TrainingArgs
|
|
from modelscope.metainfo import Trainers
|
|
from modelscope.trainers.builder import build_trainer
|
|
|
|
|
|
@dataclass(init=False)
|
|
class ImageClassificationTrainingArgs(TrainingArgs):
|
|
num_classes: int = field(
|
|
default=None,
|
|
metadata={
|
|
'cfg_node': [
|
|
'model.mm_model.head.num_classes',
|
|
'model.mm_model.train_cfg.augments.0.num_classes',
|
|
'model.mm_model.train_cfg.augments.1.num_classes'
|
|
],
|
|
'help':
|
|
'number of classes',
|
|
})
|
|
|
|
topk: tuple = field(
|
|
default=None,
|
|
metadata={
|
|
'cfg_node': [
|
|
'train.evaluation.metric_options.topk',
|
|
'evaluation.metric_options.topk'
|
|
],
|
|
'help':
|
|
'evaluation using topk, tuple format, eg (1,), (1,5)',
|
|
})
|
|
|
|
warmup_iters: str = field(
|
|
default=None,
|
|
metadata={
|
|
'cfg_node': 'train.lr_config.warmup_iters',
|
|
'help': 'The warmup iters',
|
|
})
|
|
|
|
|
|
def create_dataset(name, split):
|
|
namespace, dataset_name = name.split('/')
|
|
return MsDataset.load(
|
|
dataset_name, namespace=namespace, subset_name='default', split=split)
|
|
|
|
|
|
training_args = ImageClassificationTrainingArgs(
|
|
model='damo/cv_vit-base_image-classification_ImageNet-labels',
|
|
max_epochs=1,
|
|
lr=1e-4,
|
|
optimizer='AdamW',
|
|
warmup_iters=1,
|
|
topk=(1, )).parse_cli()
|
|
config, args = training_args.to_config()
|
|
|
|
|
|
def cfg_modify_fn(cfg):
|
|
if args.use_model_config:
|
|
cfg.merge_from_dict(config)
|
|
else:
|
|
cfg = config
|
|
return cfg
|
|
|
|
|
|
def train():
|
|
train_dataset = create_dataset(
|
|
training_args.train_dataset_name, split=training_args.train_split)
|
|
val_dataset = create_dataset(
|
|
training_args.val_dataset_name, split=training_args.val_split)
|
|
|
|
kwargs = dict(
|
|
model=args.model, # model id
|
|
train_dataset=train_dataset, # training dataset
|
|
eval_dataset=val_dataset, # validation dataset
|
|
cfg_modify_fn=cfg_modify_fn # callback to modify configuration
|
|
)
|
|
|
|
# in distributed training, specify pytorch launcher
|
|
if 'MASTER_ADDR' in os.environ:
|
|
kwargs['launcher'] = 'pytorch'
|
|
|
|
trainer = build_trainer(
|
|
name=Trainers.image_classification, default_args=kwargs)
|
|
# start to train
|
|
trainer.train()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
train()
|