chemprop/cli/hpopt.py

import json
import logging
import sys
from argparse import ArgumentParser, Namespace
from copy import deepcopy
from pathlib import Path
import torch
from lightning import pytorch as pl
from lightning.pytorch.callbacks import EarlyStopping

from chemprop.cli.common import add_common_args, process_common_args, validate_common_args
from chemprop.cli.train import (
    add_train_args,
    build_datasets,
    build_model,
    build_splits,
    normalize_inputs,
    process_train_args,
    validate_train_args,
)
from chemprop.cli.utils.command import Subcommand
from chemprop.data import build_dataloader
from chemprop.featurizers import MoleculeFeaturizerRegistry
from chemprop.nn import AggregationRegistry
from chemprop.nn.transforms import UnscaleTransform
from chemprop.nn.utils import Activation
from chemprop.utils import Factory

NO_RAY = False
DEFAULT_SEARCH_SPACE = {}
try:
    import ray
    from ray import tune
    from ray.train import CheckpointConfig, RunConfig, ScalingConfig
    from ray.train.lightning import (
        RayDDPStrategy,
        RayLightningEnvironment,
        RayTrainReportCallback,
        prepare_trainer,
    )
    from ray.train.torch import TorchTrainer
    from ray.tune.schedulers import ASHAScheduler

    DEFAULT_SEARCH_SPACE = {
        "activation": tune.choice(categories=list(Activation.keys())),
        "aggregation": tune.choice(categories=list(AggregationRegistry.keys())),
        "aggregation_norm": tune.quniform(lower=1, upper=200, q=1),
        "batch_size": tune.choice([16, 32, 64, 128, 256]),
        "depth": tune.qrandint(lower=2, upper=6, q=1),
        "dropout": tune.choice([tune.choice([0.0]), tune.quniform(lower=0.05, upper=0.4, q=0.05)]),
        "ffn_hidden_dim": tune.qrandint(lower=300, upper=2400, q=100),
        "ffn_num_layers": tune.qrandint(lower=1, upper=3, q=1),
        "final_lr_ratio": tune.loguniform(lower=1e-2, upper=1),
        "message_hidden_dim": tune.qrandint(lower=300, upper=2400, q=100),
        "init_lr_ratio": tune.loguniform(lower=1e-2, upper=1),
        "max_lr": tune.loguniform(lower=1e-4, upper=1e-2),
        "warmup_epochs": None,
    }
except ImportError:
    NO_RAY = True

NO_HYPEROPT = False
try:
    from ray.tune.search.hyperopt import HyperOptSearch
except ImportError:
    NO_HYPEROPT = True

# NO_OPTUNA = False
# try:
#     from ray.tune.search.optuna import OptunaSearch
# except ImportError:
#     NO_OPTUNA = True


logger = logging.getLogger(__name__)

SEARCH_SPACE = DEFAULT_SEARCH_SPACE

SEARCH_PARAM_KEYWORDS_MAP = {
    "basic": ["depth", "ffn_num_layers", "dropout", "ffn_hidden_dim", "message_hidden_dim"],
    "learning_rate": ["max_lr", "init_lr_ratio", "final_lr_ratio", "warmup_epochs"],
    "all": list(DEFAULT_SEARCH_SPACE.keys()),
}


class HpoptSubcommand(Subcommand):
    COMMAND = "hpopt"
    HELP = "perform hyperparameter optimization on the given task"

    @classmethod
    def add_args(cls, parser: ArgumentParser) -> ArgumentParser:
        parser = add_common_args(parser)
        parser = add_train_args(parser)
        return add_hpopt_args(parser)

    @classmethod
    def func(cls, args: Namespace):
        args = process_common_args(args)
        args = process_train_args(args)
        args = process_hpopt_args(args)
        validate_common_args(args)
        validate_train_args(args)
        main(args)


def add_hpopt_args(parser: ArgumentParser) -> ArgumentParser:
    hpopt_args = parser.add_argument_group("Chemprop hyperparameter optimization arguments")

    hpopt_args.add_argument(
        "--search-parameter-keywords",
        type=str,
        nargs="+",
        default=["basic"],
        help=f"""The model parameters over which to search for an optimal hyperparameter configuration.
    Some options are bundles of parameters or otherwise special parameter operations.

    Special keywords:
        basic - the default set of hyperparameters for search: depth, ffn_num_layers, dropout, message_hidden_dim, and ffn_hidden_dim.
        learning_rate - search for max_lr, init_lr_ratio, final_lr_ratio, and warmup_epochs. The search for init_lr and final_lr values
            are defined as fractions of the max_lr value. The search for warmup_epochs is as a fraction of the total epochs used.
        all - include search for all 13 inidividual keyword options

    Individual supported parameters:
        {list(DEFAULT_SEARCH_SPACE.keys())}
    """,
    )

    hpopt_args.add_argument(
        "--hpopt-save-dir",
        type=Path,
        help="Directory to save the hyperparameter optimization results",
    )

    raytune_args = parser.add_argument_group("Ray Tune arguments")

    raytune_args.add_argument(
        "--raytune-num-samples",
        type=int,
        default=10,
        help="Passed directly to Ray Tune TuneConfig to control number of trials to run",
    )

    raytune_args.add_argument(
        "--raytune-search-algorithm",
        choices=["random", "hyperopt"],  # , "optuna"],
        default="hyperopt",
        help="Passed to Ray Tune TuneConfig to control search algorithm",
    )

    raytune_args.add_argument(
        "--raytune-num-workers",
        type=int,
        default=1,
        help="Passed directly to Ray Tune ScalingConfig to control number of workers to use",
    )

    raytune_args.add_argument(
        "--raytune-use-gpu",
        action="store_true",
        help="Passed directly to Ray Tune ScalingConfig to control whether to use GPUs",
    )

    raytune_args.add_argument(
        "--raytune-num-checkpoints-to-keep",
        type=int,
        default=1,
        help="Passed directly to Ray Tune CheckpointConfig to control number of checkpoints to keep",
    )

    raytune_args.add_argument(
        "--raytune-grace-period",
        type=int,
        default=10,
        help="Passed directly to Ray Tune ASHAScheduler to control grace period",
    )

    raytune_args.add_argument(
        "--raytune-reduction-factor",
        type=int,
        default=2,
        help="Passed directly to Ray Tune ASHAScheduler to control reduction factor",
    )

    hyperopt_args = parser.add_argument_group("Hyperopt arguments")

    hyperopt_args.add_argument(
        "--hyperopt-n-initial-points",
        type=int,
        default=20,
        help="Passed directly to HyperOptSearch to control number of initial points to sample",
    )

    hyperopt_args.add_argument(
        "--hyperopt-random-state-seed",
        type=int,
        default=None,
        help="Passed directly to HyperOptSearch to control random state seed",
    )

    return parser


def process_hpopt_args(args: Namespace) -> Namespace:
    if args.hpopt_save_dir is None:
        args.hpopt_save_dir = Path(f"chemprop_hpopt/{args.data_path.stem}")

    args.hpopt_save_dir.mkdir(exist_ok=True, parents=True)

    search_parameters = set()

    for keyword in args.search_parameter_keywords:
        if keyword not in SEARCH_PARAM_KEYWORDS_MAP and keyword not in SEARCH_SPACE:
            raise ValueError(
                f"Search parameter keyword: {keyword} not in available options: {list(SEARCH_PARAM_KEYWORDS_MAP.keys()) + list(SEARCH_SPACE.keys())}."
            )

        search_parameters.update(
            SEARCH_PARAM_KEYWORDS_MAP[keyword]
            if keyword in SEARCH_PARAM_KEYWORDS_MAP
            else [keyword]
        )

    args.search_parameter_keywords = list(search_parameters)

    return args


def build_search_space(search_parameters: list[str], train_epochs: int) -> dict:
    if "warmup_epochs" in search_parameters and SEARCH_SPACE.get("warmup_epochs", None) is None:
        SEARCH_SPACE["warmup_epochs"] = tune.qrandint(lower=1, upper=train_epochs // 2, q=1)

    return {param: SEARCH_SPACE[param] for param in search_parameters}


def update_args_with_config(args: Namespace, config: dict) -> Namespace:
    args = deepcopy(args)

    for key, value in config.items():
        match key:
            case "final_lr_ratio":
                setattr(args, "final_lr", value * args.max_lr)

            case "init_lr_ratio":
                setattr(args, "init_lr", value * args.max_lr)

            case _:
                assert key in args, f"Key: {key} not found in args."
                setattr(args, key, value)

    return args


def train_model(config, args, train_dset, val_dset, logger, output_transform, input_transforms):
    update_args_with_config(args, config)

    train_loader = build_dataloader(
        train_dset, args.batch_size, args.num_workers, seed=args.data_seed
    )
    val_loader = build_dataloader(val_dset, args.batch_size, args.num_workers, shuffle=False)

    seed = args.pytorch_seed if args.pytorch_seed is not None else torch.seed()

    torch.manual_seed(seed)

    model = build_model(args, train_loader.dataset, output_transform, input_transforms)
    logger.info(model)

    monitor_mode = "min" if model.metrics[0].minimize else "max"
    logger.debug(f"Evaluation metric: '{model.metrics[0].alias}', mode: '{monitor_mode}'")

    patience = args.patience if args.patience is not None else args.epochs
    early_stopping = EarlyStopping("val_loss", patience=patience, mode=monitor_mode)

    trainer = pl.Trainer(
        accelerator=args.accelerator,
        devices=args.devices,
        max_epochs=args.epochs,
        gradient_clip_val=args.grad_clip,
        strategy=RayDDPStrategy(find_unused_parameters=True),
        callbacks=[RayTrainReportCallback(), early_stopping],
        plugins=[RayLightningEnvironment()],
        deterministic=args.pytorch_seed is not None,
    )
    trainer = prepare_trainer(trainer)
    trainer.fit(model, train_loader, val_loader)


def tune_model(
    args, train_dset, val_dset, logger, monitor_mode, output_transform, input_transforms
):
    scheduler = ASHAScheduler(
        max_t=args.epochs,
        grace_period=min(args.raytune_grace_period, args.epochs),
        reduction_factor=args.raytune_reduction_factor,
    )

    scaling_config = ScalingConfig(
        num_workers=args.raytune_num_workers, use_gpu=args.raytune_use_gpu
    )

    checkpoint_config = CheckpointConfig(
        num_to_keep=args.raytune_num_checkpoints_to_keep,
        checkpoint_score_attribute="val_loss",
        checkpoint_score_order=monitor_mode,
    )

    run_config = RunConfig(
        checkpoint_config=checkpoint_config,
        storage_path=args.hpopt_save_dir.absolute() / "ray_results",
    )

    ray_trainer = TorchTrainer(
        lambda config: train_model(
            config, args, train_dset, val_dset, logger, output_transform, input_transforms
        ),
        scaling_config=scaling_config,
        run_config=run_config,
    )

    match args.raytune_search_algorithm:
        case "random":
            search_alg = None
        case "hyperopt":
            if NO_HYPEROPT:
                raise ImportError(
                    "HyperOptSearch requires hyperopt to be installed. Use 'pip -U install hyperopt' to install."
                )

            search_alg = HyperOptSearch(
                n_initial_points=args.hyperopt_n_initial_points,
                random_state_seed=args.hyperopt_random_state_seed,
            )
        # case "optuna":
        #     if NO_OPTUNA:
        #         raise ImportError(
        #             "OptunaSearch requires optuna to be installed. Use 'pip -U install optuna' to install."
        #         )

        #     search_alg = OptunaSearch()

    tune_config = tune.TuneConfig(
        metric="val_loss",
        mode=monitor_mode,
        num_samples=args.raytune_num_samples,
        scheduler=scheduler,
        search_alg=search_alg,
    )

    tuner = tune.Tuner(
        ray_trainer,
        param_space={
            "train_loop_config": build_search_space(args.search_parameter_keywords, args.epochs)
        },
        tune_config=tune_config,
    )

    return tuner.fit()


def main(args: Namespace):
    if NO_RAY:
        raise ImportError(
            "Ray Tune requires ray to be installed. Use 'pip -U install ray[tune]' to install."
        )

    format_kwargs = dict(
        no_header_row=args.no_header_row,
        smiles_cols=args.smiles_columns,
        rxn_cols=args.reaction_columns,
        target_cols=args.target_columns,
        ignore_cols=args.ignore_columns,
        splits_col=args.splits_column,
        weight_col=args.weight_column,
        bounded=args.loss_function is not None and "bounded" in args.loss_function,
    )

    if args.features_generators is not None:
        # TODO: MorganFeaturizers take radius, length, and include_chirality as arguements. Should we expose these through the CLI?
        features_generators = [
            Factory.build(MoleculeFeaturizerRegistry[features_generator])
            for features_generator in args.features_generators
        ]
    else:
        features_generators = None

    featurization_kwargs = dict(
        features_generators=features_generators, keep_h=args.keep_h, add_h=args.add_h
    )

    train_data, val_data, test_data = build_splits(args, format_kwargs, featurization_kwargs)
    train_dset, val_dset, test_dset = build_datasets(args, train_data[0], val_data[0], test_data[0])

    input_transforms = normalize_inputs(train_dset, val_dset, args)

    if "regression" in args.task_type:
        output_scaler = train_dset.normalize_targets()
        val_dset.normalize_targets(output_scaler)
        logger.info(f"Train data: mean = {output_scaler.mean_} | std = {output_scaler.scale_}")
        output_transform = UnscaleTransform.from_standard_scaler(output_scaler)
    else:
        output_transform = None

    train_loader = build_dataloader(
        train_dset, args.batch_size, args.num_workers, seed=args.data_seed
    )

    model = build_model(args, train_loader.dataset, output_transform, input_transforms)
    monitor_mode = "min" if model.metrics[0].minimize else "max"

    results = tune_model(
        args, train_dset, val_dset, logger, monitor_mode, output_transform, input_transforms
    )

    best_result = results.get_best_result()
    best_config = best_result.config
    best_checkpoint = best_result.checkpoint  # Get best trial's best checkpoint

    logger.info(f"Saving best hyperparameter parameters: {best_config}")

    with open(args.hpopt_save_dir / "best_params.json", "w") as f:
        json.dump(best_config, f, indent=4)

    logger.info(f"Saving best hyperparameter configuration checkpoint: {best_checkpoint}")

    torch.save(best_checkpoint, args.hpopt_save_dir / "best_checkpoint.ckpt")

    result_df = results.get_dataframe()

    logger.info(f"Saving hyperparameter optimization results: {result_df}")

    result_df.to_csv(args.hpopt_save_dir / "all_progress.csv", index=False)

    ray.shutdown()


if __name__ == "__main__":
    parser = ArgumentParser()
    parser = HpoptSubcommand.add_args(parser)

    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)
    args = parser.parse_args()
    HpoptSubcommand.func(args)