raytune.py

import os
import time
from pathlib import Path

import ray
from ray import air, tune
# from ray.tune.schedulers import ASHAScheduler

import train

# from ray_configs.dataset_versions import search_space, run_name  # Tune #01
# from ray_configs.loss_functions import run_name, search_space  # Tune #02
# from ray_configs.head_networks import run_name, search_space  # Tune #03
# from ray_configs.augmentations import run_name, search_space  # Tune #04
from ray_configs.longer import run_name, search_space  # Tune #04

is_linux = os.name == "posix"
ray_results_dir = Path("./ray_results").resolve().as_posix()
ray_temp_dir = "/mnt/tank/fks-js/ray_temp/" if is_linux else "~/ray_temp/"
os.environ["TUNE_RESULT_DIR"] = Path("./ray_results").resolve().as_posix()


def run_tuner():
    ray.init(_temp_dir=ray_temp_dir)

    # Running the parser once with dummy dataset path to get the defaults.
    args = train.parser.parse_args(["dummy"]).__dict__

    args.update(search_space)

    start_time = time.time()

    # START PyTorch Doc
    # num_samples = 10
    # gpus_per_trial = 1/2
    # result = tune.run(
    #     train.train,
    #     resources_per_trial={"cpu": 4, "gpu": gpus_per_trial},
    #     config=search_space,
    #     num_samples=num_samples,
    #     # scheduler=scheduler,
    #     # checkpoint_at_end=True
    # )
    # results_df = result.results_df
    # END PyTorch Doc

    # START ray tune quickstart - GRID SEARCH

    resources = {"cpu": 8, "gpu": 0.5}
    trainable = tune.with_resources(train.train, resources)
    tuner = tune.Tuner(
        trainable,
        run_config=air.RunConfig(name=run_name),
        # num_samples will repeat the entire config n times.
        # when using grid_search this will multiply the number of samples by n.
        # tune_config=tune.TuneConfig(num_samples=2),
        param_space=args,
    )
    results = tuner.fit()
    # END ray tune quickstart

    results_df = results.get_dataframe()
    results_df.to_pickle(f"tune_result_{start_time}.pkl")
    # results_df.to_excel(f"tune_result_{start_time}.xlsx")  # requires openpyxl


if __name__ == "__main__":
    run_tuner()