modelscope/examples/pytorch/finetune_image_classification.py

import os
from modelscope.metainfo import Trainers
from modelscope.msdatasets.ms_dataset import MsDataset
from modelscope.trainers.builder import build_trainer
from modelscope.trainers.training_args import ArgAttr, CliArgumentParser, training_args


def define_parser():
    training_args.num_classes = ArgAttr(cfg_node_name=['model.mm_model.head.num_classes',
                                                       'model.mm_model.train_cfg.augments.0.num_classes',
                                                       'model.mm_model.train_cfg.augments.1.num_classes'],
                                        type=int, help='number of classes')

    training_args.train_batch_size.default = 16
    training_args.train_data_worker.default = 1
    training_args.max_epochs.default = 1
    training_args.optimizer.default = 'AdamW'
    training_args.lr.default = 1e-4
    training_args.warmup_iters = ArgAttr('train.lr_config.warmup_iters', type=int, default=1, help='number of warmup epochs')
    training_args.topk = ArgAttr(cfg_node_name=['train.evaluation.metric_options.topk',
                                                'evaluation.metric_options.topk'],
                                 default=(1,), help='evaluation using topk, tuple format, eg (1,), (1,5)')

    training_args.train_data = ArgAttr(type=str, default='tany0699/cats_and_dogs', help='train dataset')
    training_args.validation_data = ArgAttr(type=str, default='tany0699/cats_and_dogs', help='validation dataset')
    training_args.model_id = ArgAttr(type=str, default='damo/cv_vit-base_image-classification_ImageNet-labels', help='model name')

    parser = CliArgumentParser(training_args)
    return parser


def create_dataset(name, split):
    namespace, dataset_name = name.split('/')
    return MsDataset.load(dataset_name, namespace=namespace,
                          subset_name='default',
                          split=split)


def train(parser):
    cfg_dict = parser.get_cfg_dict()
    args = parser.args
    train_dataset = create_dataset(args.train_data, split='train')
    val_dataset = create_dataset(args.validation_data, split='validation')

    def cfg_modify_fn(cfg):
        cfg.merge_from_dict(cfg_dict)
        return cfg

    kwargs = dict(
        model=args.model_id,          # 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__':
    parser = define_parser()
    train(parser)