tune

main(cfg)

Main entry point for training.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	DictConfig configuration composed by Hydra.	required

Returns: Optional[float] with optimized metric value.

Source code in meds_torch/tune.py

@hydra.main(version_base="1.3", config_path=str(config_yaml.parent.resolve()), config_name=config_yaml.stem)
def main(cfg: DictConfig) -> float | None:
    """Main entry point for training.

    Args:
        cfg: DictConfig configuration composed by Hydra.
    Returns:
        Optional[float] with optimized metric value.
    """
    os.environ["RAY_memory_monitor_refresh_ms"] = cfg.ray_memory_monitor_refresh_ms
    # apply extra utilities
    configure_logging(cfg)

    if cfg.best_config_path:
        if not Path(cfg.best_config_path).exists():
            raise FileNotFoundError(f"Best config file not found at {cfg.best_config_path}")
        logger.info(f"Loading best tuning config from {cfg.best_config_path}")
        with open(Path(cfg.best_config_path)) as config_json:
            best_config = json.load(config_json)["train_loop_config"]
        for key, value in best_config.items():
            OmegaConf.update(cfg, key, value, merge=False)
    else:
        logger.info("No best config provided")

    analysis, best_trial = ray_tune_runner(cfg, train_func=train_func)
    ray.shutdown()

    # return tune results
    results_df = pl.from_dataframe(
        analysis.get_dataframe(
            filter_metric=cfg.hparams_search.optimized_metric, filter_mode=cfg.hparams_search.direction
        )
    )
    results_df.write_parquet(cfg.paths.time_output_dir / "sweep_results.parquet")

    best_model_path = (
        Path(
            best_trial.get_best_checkpoint(
                metric=cfg.hparams_search.optimized_metric, mode=cfg.hparams_search.direction
            ).path
        )
        / "checkpoint.ckpt"
    )

    checkpoint_dir = cfg.paths.time_output_dir / "checkpoints"
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    shutil.copy(best_model_path, checkpoint_dir / "best_model.ckpt")

    with open(cfg.paths.time_output_dir / "best_config.json", "w") as outfile:
        json.dump(best_trial.config, outfile)

    # Generate summary of results
    summary_results_df_cols = [
        each
        for each in results_df.columns
        if each.startswith("config/")
        or each.startswith("train")
        or each.startswith("val")
        or each.startswith("test")
    ] + ["logdir", "checkpoint_dir_name"]
    summary_df = results_df[summary_results_df_cols]

    # Create a new column 'best_checkpoint_path' using the get_checkpoint_path function
    log_dir_index = summary_df.columns.index("logdir")
    ckpt_dir_name_index = summary_df.columns.index("checkpoint_dir_name")
    checkpoint_paths = summary_df.map_rows(
        lambda x: get_checkpoint_path(x[log_dir_index], x[ckpt_dir_name_index], cfg.paths.time_output_dir)
    )
    summary_df = summary_df.with_columns(best_checkpoint_path=checkpoint_paths.to_series())

    if cfg.get("test"):
        logger.info("Computing Test Results")
        test_results = []
        with open_dict(cfg):
            del cfg.trainer.strategy
            del cfg.callbacks
            cfg.trainer.devices = cfg.test_devices
        datamodule = hydra.utils.instantiate(cfg.data)
        for ckpt_path in summary_df["best_checkpoint_path"].to_list():
            with open_dict(cfg):
                cfg.ckpt_path = ckpt_path
            result, _ = evaluate(cfg, datamodule=datamodule)
            test_results.append(result)
        results = {key: [result[key] for result in test_results] for key in test_results[0].keys()}
        for key, values in results.items():
            summary_df = summary_df.with_columns(pl.Series(values).alias(key))

    logger.info(summary_df)
    summary_df.write_parquet(cfg.paths.time_output_dir / "sweep_results_summary.parquet")

    return best_trial