(cfg: DictConfig)
| 40 | |
| 41 | @hydra.main(config_path="configs", config_name="config", version_base=None) |
| 42 | def main(cfg: DictConfig): |
| 43 | seed_everything(2024) |
| 44 | |
| 45 | """ Check config """ |
| 46 | check_config(cfg) |
| 47 | |
| 48 | """ Setup Logging """ |
| 49 | now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") |
| 50 | postfix = f"{cfg.slurm_id}_{now}" if exists(cfg.slurm_id) else now |
| 51 | exp_name = f"{cfg.name}_{postfix}" if exists(cfg.name) else postfix |
| 52 | log_dir = os.path.join("logs", exp_name) |
| 53 | ckpt_dir = os.path.join(log_dir, "checkpoints") |
| 54 | |
| 55 | # setup loggers |
| 56 | use_wandb_logging = cfg.use_wandb or cfg.use_wandb_offline |
| 57 | if use_wandb_logging: |
| 58 | usr_name = os.environ.get('USER', os.environ.get('USERNAME')) |
| 59 | mode = "offline" if cfg.use_wandb_offline else "online" |
| 60 | online_logger = WandbLogger( |
| 61 | dir=log_dir, |
| 62 | save_dir=log_dir, |
| 63 | name=exp_name, |
| 64 | project="lora-fm", |
| 65 | tags=[usr_name, *cfg.get("tags", [])], |
| 66 | config=OmegaConf.to_object(cfg), |
| 67 | mode=mode, |
| 68 | group="DDP" |
| 69 | ) |
| 70 | else: |
| 71 | online_logger = TensorBoardLogger( |
| 72 | save_dir=log_dir, |
| 73 | name="", |
| 74 | version="", |
| 75 | log_graph=False, |
| 76 | default_hp_metric=False, |
| 77 | ) |
| 78 | csv_logger = CSVLogger( |
| 79 | log_dir, |
| 80 | name="", |
| 81 | version="", |
| 82 | prefix="", |
| 83 | flush_logs_every_n_steps=500 |
| 84 | ) |
| 85 | csv_logger.log_hyperparams(OmegaConf.to_container(cfg)) |
| 86 | logger = [online_logger, csv_logger] |
| 87 | |
| 88 | """ Setup dataloader """ |
| 89 | data = instantiate_from_config(cfg.data) |
| 90 | |
| 91 | """ Setup model """ |
| 92 | trainer_module_params = dict( |
| 93 | # model |
| 94 | fm_cfg = cfg.model.fm_cfg, |
| 95 | noising_step = cfg.model.noising_step, |
| 96 | start_from_noise = cfg.model.start_from_noise, |
| 97 | # first stage |
| 98 | first_stage = cfg.autoencoder, |
| 99 | # lora |
no test coverage detected