<reponame>semink/TCN

from ray import tune
from ray.tune import CLIReporter

from ray.tune.schedulers import PopulationBasedTraining

from pytorch_lightning import Trainer
from argparse import ArgumentParser
from train import train_model
from model import TCNPredictor


def tune_pbt(args):
    config = {
        "lr": 1e-3,
        "batch_size": 8,
        "dropout": 0.3,
        "weight_decay": 0.01,
        "N": tune.choice([2, 4, 8]),
        "hidden_features": tune.choice([4, 8, 16, 32]),
        "shid_features": tune.choice([4, 8, 16, 32]),
        "nhid": tune.choice([1, 3, 5]),
        "kernel_size": tune.choice([3, 5, 7]),
    }

    scheduler = PopulationBasedTraining(
        perturbation_interval=4,
        hyperparam_mutations={
            "lr": tune.loguniform(1e-4, 1e-1),
            "batch_size": [8, 16],
            "dropout": tune.uniform(0, 0.99),
            "weight_decay": tune.loguniform(1e-5, 1e-1)
        })

    reporter = CLIReporter(
        parameter_columns=[c for c in config],
        metric_columns=["loss", "training_iteration"],
        max_report_frequency=20)

    analysis = tune.run(
        tune.with_parameters(train_model, args=args, tuning=True),
        resources_per_trial={"cpu": 4, "gpu": 1},
        metric="loss",
        mode="min",
        config=config,
        num_samples=args.ray_num_samples,
        scheduler=scheduler,
        progress_reporter=reporter, resume=args.resume,
        name=f"tune_TCN_pbt{args.expid}", verbose=2)
    best_trial = analysis.get_best_trial(metric='loss', mode='min')
    best_model_checkpoint = analysis.get_best_checkpoint(
        best_trial, metric='loss')
    print("Best hyperparameters found were: ", analysis.best_config)
    print("Best checkpoint: ", best_model_checkpoint)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser = Trainer.add_argparse_args(parser)
    # PROGRAM level args
    parser.add_argument("--conda_env", type=str, default="pytorch")
    parser.add_argument('--ray_num_samples', type=int, default=20,
                        help='number of ray tune samples')
    parser.add_argument('--expid', type=int, default=0,
                        help='ray tune experiment id')
    parser.add_argument('--resume', action='store_true',
                        help='whether to resume previous tuning process')

    # MODEL specific args
    parser = TCNPredictor.add_model_specific_args(parser)

    # DATALOADER specific args
    parser.add_argument("--seq_len", type=int, default=12)
    parser.add_argument("--num_workers", type=int, default=32)
    parser.add_argument("--weeks", type=int, default=5)

    args = parser.parse_args()
    tune_pbt(args)

from ray import tune
from ray.tune import CLIReporter

from ray.tune.schedulers import PopulationBasedTraining

from pytorch_lightning import Trainer
from argparse import ArgumentParser
from train import train_model
from model import TCNPredictor


def tune_pbt(args):
    config = {
        "lr": 1e-3,
        "batch_size": 8,
        "dropout": 0.3,
        "weight_decay": 0.01,
        "N": tune.choice([2, 4, 8]),
        "hidden_features": tune.choice([4, 8, 16, 32]),
        "shid_features": tune.choice([4, 8, 16, 32]),
        "nhid": tune.choice([1, 3, 5]),
        "kernel_size": tune.choice([3, 5, 7]),
    }

    scheduler = PopulationBasedTraining(
        perturbation_interval=4,
        hyperparam_mutations={
            "lr": tune.loguniform(1e-4, 1e-1),
            "batch_size": [8, 16],
            "dropout": tune.uniform(0, 0.99),
            "weight_decay": tune.loguniform(1e-5, 1e-1)
        })

    reporter = CLIReporter(
        parameter_columns=[c for c in config],
        metric_columns=["loss", "training_iteration"],
        max_report_frequency=20)

    analysis = tune.run(
        tune.with_parameters(train_model, args=args, tuning=True),
        resources_per_trial={"cpu": 4, "gpu": 1},
        metric="loss",
        mode="min",
        config=config,
        num_samples=args.ray_num_samples,
        scheduler=scheduler,
        progress_reporter=reporter, resume=args.resume,
        name=f"tune_TCN_pbt{args.expid}", verbose=2)
    best_trial = analysis.get_best_trial(metric='loss', mode='min')
    best_model_checkpoint = analysis.get_best_checkpoint(
        best_trial, metric='loss')
    print("Best hyperparameters found were: ", analysis.best_config)
    print("Best checkpoint: ", best_model_checkpoint)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser = Trainer.add_argparse_args(parser)
    # PROGRAM level args
    parser.add_argument("--conda_env", type=str, default="pytorch")
    parser.add_argument('--ray_num_samples', type=int, default=20,
                        help='number of ray tune samples')
    parser.add_argument('--expid', type=int, default=0,
                        help='ray tune experiment id')
    parser.add_argument('--resume', action='store_true',
                        help='whether to resume previous tuning process')

    # MODEL specific args
    parser = TCNPredictor.add_model_specific_args(parser)

    # DATALOADER specific args
    parser.add_argument("--seq_len", type=int, default=12)
    parser.add_argument("--num_workers", type=int, default=32)
    parser.add_argument("--weeks", type=int, default=5)

    args = parser.parse_args()
    tune_pbt(args)
# PROGRAM level args # MODEL specific args # DATALOADER specific args