Function train

classification/main_accelerate.py:301–340

(config, accelerator: Accelerator)

def train(config, accelerator: Accelerator):
    _, _, _, training_dataloader, validate_dataloader, _, mixup_fn = build_loader2(config)
    model = build_model(config)
    optimizer = build_optimizer(config, model)
    criterion = build_criterion(config)

    model, optimizer, training_dataloader, validate_dataloader = accelerator.prepare(
        model, optimizer, training_dataloader, validate_dataloader)

    effective_update_steps_per_epoch = len(training_dataloader) // config.TRAIN.ACCUMULATION_STEPS
    lr_scheduler = build_scheduler(config, optimizer, effective_update_steps_per_epoch)

    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!')

    max_acc = load_checkpoint_if_needed(accelerator, config, lr_scheduler)

    logger.info(f'Created model:{config.MODEL.TYPE}/{config.MODEL.NAME}')
    logger.info(str(model))
    logger.info('Effective Optimizer Steps: {}'.format(effective_update_steps_per_epoch))
    logger.info('Start training')
    logger.info('Max accuracy: {}'.format(max_acc))
    log_model_statistic(accelerator.unwrap_model(model))

    for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS):
        train_epoch(model=model, optimizer=optimizer, data_loader=training_dataloader,
                    scheduler=lr_scheduler, criterion=criterion, mixup_fn=mixup_fn,
                    accelerator=accelerator, epoch=epoch, config=config)
        acc = eval_epoch(config=config, data_loader=validate_dataloader, model=model,
                         accelerator=accelerator)

        accelerator.wait_for_everyone()
        if acc > max_acc:
            max_acc = acc
            save_checkpoint(os.path.join(config.OUTPUT, 'best'), accelerator, epoch, max_acc, config, lr_scheduler)
        logger.info(f'Max Acc@1 {max_acc:.3f}')
        save_checkpoint(os.path.join(config.OUTPUT, 'last'), accelerator, epoch, max_acc, config, lr_scheduler)


def main():