Function main

classification/main_accelerate.py:343–373

()

def main():
    args, config = parse_option()
    os.makedirs(config.OUTPUT, exist_ok=True)
    logging.basicConfig(
        format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
        datefmt='%m/%d/%Y %H:%M:%S',
        filename=os.path.join(config.OUTPUT, 'run.log'),
        level=logging.INFO,
    )

    loggers = ['tensorboard']
    accelerator = Accelerator(
        log_with=loggers,
        project_dir=config.OUTPUT,
        gradient_accumulation_steps=config.TRAIN.ACCUMULATION_STEPS,
        # When use deepspeed, you could not comment this out
        # even if you set loss scale to 1.0 in deepspeed config.
        kwargs_handlers=[GradScalerKwargs(enabled=not args.disable_grad_scalar)],
    )
    logger.info(accelerator.state, main_process_only=False)

    scale_learning_rate(config, accelerator.num_processes)
    seed_everything(config.SEED, accelerator.process_index)
    save_config(config)

    logger.info(config.dump())

    if config.EVAL_MODE:
        eval(config, accelerator)
    else:
        train(config, accelerator)


if __name__ == '__main__':