Function train_worker

tests/trainers/test_team_transfer_trainer.py:25–80

(device_id)

def train_worker(device_id):
    model_id = 'damo/multi-modal_team-vit-large-patch14_multi-modal-similarity'
    ckpt_dir = './ckpt'
    os.makedirs(ckpt_dir, exist_ok=True)
    # Use epoch=1 for faster training here
    cfg = Config({
        'framework': 'pytorch',
        'task': 'multi-modal-similarity',
        'pipeline': {
            'type': 'multi-modal-similarity'
        },
        'model': {
            'type': 'team-multi-modal-similarity'
        },
        'dataset': {
            'name': 'Caltech101',
            'class_num': 101
        },
        'preprocessor': {},
        'train': {
            'epoch': 1,
            'batch_size': 32,
            'ckpt_dir': ckpt_dir
        },
        'evaluation': {
            'batch_size': 64
        }
    })
    cfg_file = '{}/{}'.format(ckpt_dir, ModelFile.CONFIGURATION)
    cfg.dump(cfg_file)

    train_dataset = MsDataset.load(
        cfg.dataset.name,
        namespace='modelscope',
        split='train',
        download_mode=DownloadMode.FORCE_REDOWNLOAD).to_hf_dataset()
    train_dataset = train_dataset.with_transform(train_mapping)
    val_dataset = MsDataset.load(
        cfg.dataset.name,
        namespace='modelscope',
        split='validation',
        download_mode=DownloadMode.FORCE_REDOWNLOAD).to_hf_dataset()
    val_dataset = val_dataset.with_transform(val_mapping)

    default_args = dict(
        cfg_file=cfg_file,
        model=model_id,
        device_id=device_id,
        data_collator=collate_fn,
        train_dataset=train_dataset,
        val_dataset=val_dataset)

    trainer = build_trainer(
        name=Trainers.image_classification_team, default_args=default_args)
    trainer.train()
    trainer.evaluate()