tune_complete_xgb_wb.py

"""
tune_complete_xgb_wb.py
====================================
XGBoost tune script for Complete case on WanDB.
"""

from utility.config import load_config
from xgboost import XGBClassifier
from tools import data_loader
import paths as pt
from sklearn.model_selection import cross_validate, StratifiedKFold
import numpy as np
import os

os.environ["WANDB_SILENT"] = "true"
import wandb

sweep_config = {
    "method": "random", # try grid or random
    "metric": {
      "name": "accuracy",
      "goal": "maximize"
    },
    "parameters": {
        "n_estimators": {
            "values": [50, 100, 200, 400]
        },
        "booster": {
            "values": ["gbtree"]
        },
        "max_depth": {
            "values": [3, 4, 5, 6, 7, 8, 9, 10]
        },
        "gamma": {
            "values": [0, 1, 5]
        },
        "colsample_bytree": {
            "values": [0.3, 0.5, 0.8, 1]
        },
        "min_child_weight": {
            "values": [int(x) for x in np.linspace(1, 10, 10, endpoint=True)]
        },
        "reg_alpha": {
            "values": [int(x) for x in np.linspace(10, 180, 20, endpoint=True)]
        },
        "reg_lambda": {
            "values": [float(x) for x in np.linspace(0, 0.9, 10, endpoint=True)]
        },
        "learning_rate": {
            "values": [0.01, 0.1, 0.05, 0.2]
        },
        "subsample": {
            "values": [0.8, 1]
        },
        "scale_pos_weight": {
            "values": [False, True]
        },
    }
}

def main():
    sweep_id = wandb.sweep(sweep_config,
                           project="air-complete-xgb")
    wandb.agent(sweep_id, train_model, count=100)

def train_model():
    config_defaults = {
        "n_estimators": 200,
        "booster": "gbtree",
        "max_depth": 3,
        "gamma": 0,
        "colsample_bytree": 1,
        "min_child_weight": 1,
        "reg_alpha": 0,
        "reg_lambda": 1,
        "learning_rate": 0.3,
        "subsample": 1,
        'scale_pos_weight': False,
        "seed": 0,
        "test_size": 0.3,
    }

    # Initialize a new wandb run
    wandb.init(config=config_defaults)

    # Config is a variable that holds and saves hyperparameters and inputs
    config = wandb.config

    # Load data
    settings = load_config(pt.CONFIGS_DIR, "complete.yaml")
    dl = data_loader.CompleteDataLoader(pt.PROCESSED_DATA_DIR,
                                        "complete_emb.csv",
                                        settings).load_data()
    X, y = dl.get_data()

    # Make model
    if config.scale_pos_weight:
        neg, pos = np.bincount(y)
        scale_pos_weight = neg / pos
    else:
        scale_pos_weight = None
    params = {"n_estimators": config.n_estimators,
              "booster": config.booster,
              "max_depth":config.max_depth,
              "gamma":config.gamma,
              "colsample_bytree":config.colsample_bytree,
              "min_child_weight":config.min_child_weight,
              "reg_alpha":config.reg_alpha,
              "reg_lambda":config.reg_lambda,
              "learning_rate":config.learning_rate,
              "subsample":config.subsample,
              "scale_pos_weight":scale_pos_weight,
              "use_label_encoder":False,
              "eval_metric":"logloss",
              "objective":"binary:logistic",
              "random_state":0}
    model = XGBClassifier(**params)

    # Make CV
    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
    metrics = ['accuracy', 'precision', 'recall',
               'roc_auc', 'average_precision', 'f1']
    res_validate = cross_validate(model, X, y, cv=skf, scoring=metrics)

    # Evaluate performance
    accuracy = round(np.mean(res_validate[f'test_accuracy']), 3)
    precision = round(np.mean(res_validate[f'test_precision']), 3)
    recall = round(np.mean(res_validate[f'test_recall']), 3)
    roc_auc = round(np.mean(res_validate[f'test_roc_auc']), 3)
    avg_prec = round(np.mean(res_validate['test_average_precision']), 3)
    f1 = round(np.mean(res_validate['test_f1']), 3)

    # Log to wandb
    wandb.log({"accuracy": accuracy})
    wandb.log({"precision": precision})
    wandb.log({"recall": recall})
    wandb.log({"roc_auc": roc_auc})
    wandb.log({"pr_auc": avg_prec})
    wandb.log({"f1": f1})

if __name__ == "__main__":
    main()