notebooks/lightgbm_sorel20m.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# LightGBM calibration experiment on Sorel20M Dataset\n",
    "\n",
    "We compare multiple calibration algorithms with their reduced, \n",
    "class-wise and class-wise reduced counterparts.\n",
    "\n",
    "We work with a pre-trained LightGBM classifier trained on the [SOREL20M Dataset](https://github.com/sophos/SOREL-20M),\n",
    "a binary classification dataset consisting of nearly\n",
    "20 million malicious and benign portable executable\n",
    "files with pre-extracted features and metadata, and high quality labels\n",
    "\n",
    "The model achieves an accuracy of roughly 98% of the test set.\n",
    "\n",
    "Since the model's accuracy is pretty high it is, as expected, well calibrated\n",
    "(pre-calibration ECE ≈ 0.005, post-calibration ECE <= 0.002).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import types\n",
    "from typing import List\n",
    "\n",
    "import lightgbm as lgb\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from kyle.evaluation import EvalStats\n",
    "from sklearn.metrics import accuracy_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is needed for notebooks in case jupyter is started directly in the notebooks directory\n",
    "current_working_directory = Path(\".\").resolve()\n",
    "if current_working_directory.name == \"notebooks\":\n",
    "    sys.path.insert(0, os.fspath(current_working_directory.parent))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.constants import DATA_DIR, OUTPUT_DIR, RANDOM_SEED\n",
    "from src.data_and_models.sorel20m import download_sorel20m\n",
    "from src.utils import (\n",
    "    configure_plots,\n",
    "    perform_default_evaluation,\n",
    "    plot_evaluation_results_from_dataframe,\n",
    "    set_random_seed,\n",
    ")\n",
    "from src.utils.evaluation import combined_results_into_dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = OUTPUT_DIR / \"lightgbm_sorel20m\"\n",
    "output_dir.mkdir(exist_ok=True)\n",
    "output_file = output_dir / \"results.csv\"\n",
    "\n",
    "sorel20m_dir = DATA_DIR / \"sorel20m\"\n",
    "lightgbm_model_file = sorel20m_dir / \"lightgbm.model\"\n",
    "features_dir = sorel20m_dir / \"test-features\"\n",
    "features_file = features_dir / \"arr_0.npy\" / \"arr_0.npy\"\n",
    "labels_file = features_dir / \"arr_1.npy\" / \"arr_1.npy\"\n",
    "\n",
    "n_classes = 2\n",
    "classes = [\"malware\", \"benign\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_random_seed(RANDOM_SEED)\n",
    "configure_plots()\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(message)s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "download_sorel20m(sorel20m_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test = np.load(features_file, mmap_mode=\"r\")\n",
    "\n",
    "# Cannot use mmap with netcal because it does stupid things like:\n",
    "# Checking for array types this way: type(X) != np.ndarray\n",
    "# Instead of: isinstance(X, np.ndarray), which doesn't work for memory mapped arrays\n",
    "# because they are a subclass of np.ndarray\n",
    "# test_labels = np.load(labels_file, mmap_mode=\"r\")\n",
    "y_test = np.load(labels_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "values, counts = np.unique(y_test, return_counts=True)\n",
    "ax.bar(values, counts, edgecolor=\"k\", linewidth=2)\n",
    "ax.set_xticks(values)\n",
    "ax.set_xticklabels(classes, rotation=45)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "The following hack is needed because kyle expects classifiers to have the `predict_proba` method\n",
    "and to output an array of the same dimensionality as the number of classes\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_proba(self, X):\n",
    "    y = model.predict(X)\n",
    "    y = y[:, np.newaxis]\n",
    "    y = np.append(\n",
    "        np.zeros(y.shape, dtype=float),\n",
    "        y,\n",
    "        axis=1,\n",
    "    )\n",
    "    y[:, 0] = 1.0 - y[:, 1]\n",
    "    return y\n",
    "\n",
    "\n",
    "model = lgb.Booster(model_file=lightgbm_model_file)\n",
    "model.predict_proba = types.MethodType(predict_proba, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lgb.plot_tree(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating Calibration\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uncalibrated_confidences = model.predict_proba(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = np.argmax(uncalibrated_confidences, axis=1)\n",
    "model_accuracy = accuracy_score(y_test, y_pred)\n",
    "logger.info(f\"Model accuracy: {model_accuracy*100}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_stats = EvalStats(y_test, uncalibrated_confidences, bins=25)\n",
    "logger.info(f\"ECE before calibration: {eval_stats.expected_calibration_error()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recalibration\n",
    "We evaluate reduction wrappers on multiple metrics with different calibration algorithms\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Balanced\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Performing evaluation\")\n",
    "\n",
    "eval_results = perform_default_evaluation(\n",
    "    confidences=uncalibrated_confidences,\n",
    "    gt_labels=y_test,\n",
    ")\n",
    "\n",
    "results_df = combined_results_into_dataframe(\n",
    "    eval_results,\n",
    "    model_name=\"LightGBM\",\n",
    "    dataset_name=\"SOREL20M\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reduction_methods_order: List[str] = results_df[\"Reduction Method\"].unique().tolist()\n",
    "reduction_methods_order = [reduction_methods_order[0]] + sorted(\n",
    "    reduction_methods_order[1:], key=len\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save Results\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Saving results\")\n",
    "results_df.to_csv(output_file, sep=\";\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Plotting results\")\n",
    "\n",
    "results_df = results_df.query(\"(Metric != 'condition') & (Metric != 'weak_condition')\")\n",
    "\n",
    "plot_evaluation_results_from_dataframe(\n",
    "    results_df,\n",
    "    hue_order=reduction_methods_order,\n",
    "    output_file=(output_dir / \"evaluation_ECE_lightgbm_sorel20m.eps\"),\n",
    "    show=False,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}