Function train

classification/main_deepspeed.py:336–454

(config, ds_config)

def train(config, ds_config):
    # -------------- build ---------------- #

    _, dataset_val, _, data_loader_train, data_loader_val, _, mixup_fn = build_loader(config)
    model = build_model(config)
    model.cuda()

    if config.MODEL.PRETRAINED:
        load_pretrained(config, model, logger)

    logger.info(ds_config)
    model, optimizer, _, _ = deepspeed.initialize(
        config=ds_config,
        model=model,
        model_parameters=get_parameter_groups(model, config),
        dist_init_required=False,
    )

    try:
        model.register_comm_hook(state=None, hook=fp16_compress_hook)
        logger.info('using fp16_compress_hook!')
    except:
        logger.info('cannot register fp16_compress_hook!')

    model_without_ddp = model.module

    lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train))
    criterion = build_criterion(config)

    model_ema = None
    if config.TRAIN.EMA.ENABLE:
        model_ema = EMADeepspeed(model, config.TRAIN.EMA.DECAY)

    # -------------- resume ---------------- #

    max_accuracy = 0.0
    max_accuracy_ema = 0.0
    client_state = {}
    if config.MODEL.RESUME == '' and config.TRAIN.AUTO_RESUME:
        if os.path.exists(os.path.join(config.OUTPUT, 'latest')):
            config.defrost()
            config.MODEL.RESUME = config.OUTPUT
            config.freeze()
            tag = None
    elif config.MODEL.RESUME:
        config.MODEL.RESUME = os.path.dirname(config.MODEL.RESUME)
        tag = os.path.basename(config.MODEL.RESUME)
    if config.MODEL.RESUME:
        logger.info('loading checkpoint from {}'.format(config.MODEL.RESUME))
        _, client_state = model.load_checkpoint(load_dir=config.MODEL.RESUME, tag=tag)
        logger.info(f'client_state={client_state.keys()}')
        lr_scheduler.load_state_dict(client_state['custom_lr_scheduler'])
        max_accuracy = client_state['max_accuracy']

        if model_ema is not None:
            max_accuracy_ema = client_state.get('max_accuracy_ema', 0.0)
            model_ema.load_state_dict((client_state['model_ema']))