Function main

train_acc.py:113–650

(args)

@hydra.main(config_path="config", config_name="default", version_base=None)
def main(args):
    assert torch.cuda.is_available(), "Training currently requires at least one GPU."
    from accelerate.utils import AutocastKwargs

    kwargs = AutocastKwargs(enabled=True)
    # https://github.com/pytorch/pytorch/issues/40497#issuecomment-709846922
    # https://github.com/huggingface/accelerate/issues/2487#issuecomment-1969997224

    accelerator = accelerate.Accelerator(
        kwargs_handlers=[kwargs], mixed_precision=args.mixed_precision
    )
    device = accelerator.device
    accelerate.utils.set_seed(args.global_seed, device_specific=True)
    rank = accelerator.state.process_index
    logging.info(
        f"Starting rank={rank}, world_size={accelerator.state.num_processes}, accelerator.mixed_precision={accelerator.mixed_precision},device={device}."
    )
    is_multiprocess = True if accelerator.state.num_processes > 1 else False
    if accelerator.state.num_processes >= 4 * 8:
        args.data.sample_fid_n = min(args.data.sample_fid_n, 1_000)
        print(
            "decreasing sample_fid_n to 1_000 in node with >= 4*8 GPUs, an unknown bug from torchmetrics"
        )

    _fid_eval_batch_nums = args.data.sample_fid_n // (
        args.data.sample_fid_bs * accelerator.state.num_processes
    )
    assert _fid_eval_batch_nums > 0, f"{_fid_eval_batch_nums} <= 0"

    slurm_job_id = os.environ.get("SLURM_JOB_ID")
    logging.info(f"slurm_job_id: {slurm_job_id}")

    local_bs = args.data.batch_size

    train_steps = 0

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        logging.info(args)
        experiment_dir = HydraConfig.get().run.dir
        logging.info(f"Experiment directory created at {experiment_dir}")
        checkpoint_dir = (
            f"{experiment_dir}/checkpoints"  # Stores saved model checkpoints
        )
        os.makedirs(checkpoint_dir, exist_ok=True)
        logger = create_logger(rank, experiment_dir)
        logger.info(f"Experiment directory created at {experiment_dir}")

        if args.use_wandb:
            config_dict = OmegaConf.to_container(args, resolve=True)
            config_dict = {
                **config_dict,
                "experiment_dir": experiment_dir,
                "world_size": accelerator.state.num_processes,
                "local_batch_size": args.data.batch_size
                * accelerator.state.num_processes,
                "job_id": slurm_job_id,
            }