mirror of
https://github.com/modelscope/modelscope.git
synced 2025-12-23 11:39:26 +01:00
34 lines
957 B
Python
34 lines
957 B
Python
|
from dataclasses import dataclass, field
|
||
|
|
|
||
|
|
from modelscope.msdatasets import MsDataset
|
||
|
|
from modelscope.trainers import EpochBasedTrainer, build_trainer
|
||
|
|
from modelscope.trainers.training_args import TrainingArgs
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class StableDiffusionArguments(TrainingArgs):
|
||
|
|
|
||
|
|
def __call__(self, config):
|
||
|
|
config = super().__call__(config)
|
||
|
|
config.train.lr_scheduler.T_max = self.max_epochs
|
||
|
|
config.model.inference = False
|
||
|
|
return config
|
||
|
|
|
||
|
|
|
||
|
|
args = StableDiffusionArguments.from_cli(task='efficient-diffusion-tuning')
|
||
|
|
print(args)
|
||
|
|
|
||
|
|
dataset = MsDataset.load(args.dataset_name, namespace=args.namespace)
|
||
|
|
train_dataset = dataset['train']
|
||
|
|
validation_dataset = dataset['validation']
|
||
|
|
|
||
|
|
kwargs = dict(
|
||
|
|
model=args.model,
|
||
|
|
work_dir=args.work_dir,
|
||
|
|
train_dataset=train_dataset,
|
||
|
|
eval_dataset=validation_dataset,
|
||
|
|
cfg_modify_fn=args)
|
||
|
|
|
||
|
|
trainer: EpochBasedTrainer = build_trainer(name='trainer', default_args=kwargs)
|
||
|
|
trainer.train()
|