From 84203c63f0d9886356cda678ff1c0d198ece36c1 Mon Sep 17 00:00:00 2001 From: Evan Mendenhall Date: Mon, 8 Jun 2026 17:11:10 -0700 Subject: [PATCH] Add neural network tutorial for defect prediction --- .../.gitignore | 34 +++ .../README.md | 177 ++++++++++++++++ .../config/default.yaml | 42 ++++ .../data/processed/.gitkeep | 1 + .../data/raw/.gitkeep | 1 + .../defect_prediction.ipynb | 200 ++++++++++++++++++ .../requirements.txt | 5 + .../scripts/01_logistic_regression.py | 51 +++++ .../scripts/02_gradient_boosted_trees.py | 63 ++++++ .../scripts/03_compare_features.py | 69 ++++++ .../scripts/04_train_neural_network.py | 52 +++++ .../scripts/download_data.py | 37 ++++ .../src/__init__.py | 1 + .../src/config.py | 90 ++++++++ .../src/data/__init__.py | 24 +++ .../src/data/dataset.py | 125 +++++++++++ .../src/data/nasa_preprocessing.py | 98 +++++++++ .../src/model/__init__.py | 3 + .../src/model/defect_predictor.py | 42 ++++ .../src/paths.py | 21 ++ .../src/training/__init__.py | 3 + .../src/training/trainer.py | 140 ++++++++++++ 22 files changed, 1279 insertions(+) create mode 100644 defect_prediction_neural_network_tutorial/.gitignore create mode 100644 defect_prediction_neural_network_tutorial/README.md create mode 100644 defect_prediction_neural_network_tutorial/config/default.yaml create mode 100644 defect_prediction_neural_network_tutorial/data/processed/.gitkeep create mode 100644 defect_prediction_neural_network_tutorial/data/raw/.gitkeep create mode 100644 defect_prediction_neural_network_tutorial/defect_prediction.ipynb create mode 100644 defect_prediction_neural_network_tutorial/requirements.txt create mode 100644 defect_prediction_neural_network_tutorial/scripts/01_logistic_regression.py create mode 100644 defect_prediction_neural_network_tutorial/scripts/02_gradient_boosted_trees.py create mode 100644 defect_prediction_neural_network_tutorial/scripts/03_compare_features.py create mode 100644 defect_prediction_neural_network_tutorial/scripts/04_train_neural_network.py create mode 100644 defect_prediction_neural_network_tutorial/scripts/download_data.py create mode 100644 defect_prediction_neural_network_tutorial/src/__init__.py create mode 100644 defect_prediction_neural_network_tutorial/src/config.py create mode 100644 defect_prediction_neural_network_tutorial/src/data/__init__.py create mode 100644 defect_prediction_neural_network_tutorial/src/data/dataset.py create mode 100644 defect_prediction_neural_network_tutorial/src/data/nasa_preprocessing.py create mode 100644 defect_prediction_neural_network_tutorial/src/model/__init__.py create mode 100644 defect_prediction_neural_network_tutorial/src/model/defect_predictor.py create mode 100644 defect_prediction_neural_network_tutorial/src/paths.py create mode 100644 defect_prediction_neural_network_tutorial/src/training/__init__.py create mode 100644 defect_prediction_neural_network_tutorial/src/training/trainer.py diff --git a/defect_prediction_neural_network_tutorial/.gitignore b/defect_prediction_neural_network_tutorial/.gitignore new file mode 100644 index 0000000..a7342a6 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/.gitignore @@ -0,0 +1,34 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +.eggs/ +dist/ +build/ +*.egg +.venv/ +venv/ +env/ + +# ML artifacts +checkpoints/ +runs/ +logs/ +*.pt +*.pth +wandb/ + +# Data (keep structure, ignore large/local datasets) +data/raw/* +data/processed/* +!data/raw/.gitkeep +!data/processed/.gitkeep + +# IDE / OS +.idea/ +.vscode/ +.DS_Store +*.swp + +# Jupyter +.ipynb_checkpoints/ diff --git a/defect_prediction_neural_network_tutorial/README.md b/defect_prediction_neural_network_tutorial/README.md new file mode 100644 index 0000000..057d5be --- /dev/null +++ b/defect_prediction_neural_network_tutorial/README.md @@ -0,0 +1,177 @@ +# Software Defect Prediction Tutorial + +Learn to predict software defects using NASA metrics data. You will try three classical methods, pick the best input feature, then train a neural network that achieves the highest accuracy. + +This tutorial is based on [defect_prediction.ipynb](defect_prediction.ipynb) and the NASA PROMISE dataset ([JM1](http://promise.site.uottawa.ca/SERepository/datasets/jm1.arff) / [CM1](http://promise.site.uottawa.ca/SERepository/datasets/cm1.arff)). + +## What you will build + +| Step | Method | Goal | +|------|--------|------| +| 1 | Logistic regression | Notebook baseline (~73% AUC) | +| 2 | Gradient boosted trees | Stronger tabular baseline | +| 3 | Feature comparison | Swap the 6th feature (`i` → `total_Opnd`) | +| 4 | Neural network | Best result (~79% AUC on CM1) | + +## Setup + +Requires Python 3.10+. + +```bash +cd defect_prediction_neural_network_tutorial + +python3 -m venv .venv +source .venv/bin/activate # Windows: .venv\Scripts\activate +pip install -r requirements.txt + +python scripts/download_data.py # skip if data/raw/jm1.csv already exists +``` + +Downloading pytorch libraries for neural network training takes a few minutes: +2026-06-08 16:52:32.862 [info] ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 88.0/88.0 MB 477.0 kB/s 0:03:04 + +All scripts resolve paths from the project root (`src/paths.py`), so they work even if you run them from another directory: + +```bash +python /path/to/accuracy/scripts/01_logistic_regression.py +``` + +## Tutorial steps + +Run these in order. Each step prints its results to the terminal. + +### Step 1 — Logistic regression + +Uses six features from the notebook: `loc`, `d`, `locCodeAndComment`, `v(g)`, `uniq_Opnd`, `i`. + +```bash +python scripts/01_logistic_regression.py +``` + +**Expected:** AUC-ROC around **0.727** on an 80/20 split of JM1. + +### Step 2 — Gradient boosted trees + +Same features and split as Step 1, using scikit-learn's gradient boosted trees. + +```bash +python scripts/02_gradient_boosted_trees.py +``` + +**Expected:** AUC-ROC similar to or slightly above Step 1. + +### Step 3 — Pick the best 6th feature + +Five features stay fixed. This step tries every other NASA metric as the 6th input and ranks them. + +```bash +python scripts/03_compare_features.py +``` + +**Expected output (top of ranking):** + +Testing 16 options for the 6th feature slot + +Rank Feature AUC-ROC +------------------------------------ +1 total_Opnd 0.7319 <-- best +2 lOComment 0.7308 +3 branchCount 0.7308 +4 ev(g) 0.7305 +5 lOBlank 0.7305 +6 iv(g) 0.7300 +7 i 0.7296 +8 b 0.7275 +9 e 0.7273 +10 t 0.7273 +11 v 0.7273 +12 n 0.7269 +13 lOCode 0.7260 +14 total_Op 0.7221 +15 uniq_Op 0.7198 +16 l 0.7117 + +The default 6th feature `i` is good, but **`total_Opnd`** scores highest. Update the config: + +```yaml +# config/default.yaml +variant_feature: total_Opnd +``` + +### Step 4 — Train the neural network + +Trains on all of **JM1** and evaluates on **CM1** (a different NASA project). This is harder than Steps 1–3 but reflects real cross-project prediction. + +```bash +python scripts/04_train_neural_network.py +``` +Note: You can train a 92M param model in about 2 minutes on an M4 Pro Max Chip, and even if it is slower, watching the process is both fun and educational! No external GPUs needed, trains locally. + +**Expected:** ~92M parameter model, 20 epochs. Best validation AUC around **0.79** with `total_Opnd`: + +``` +Epoch 20/20 | val_auc=0.7887 +``` + +With the default feature `i`, expect val AUC around **0.77**. + +## Project layout + +``` +accuracy/ +├── README.md +├── defect_prediction.ipynb # original notebook exercise +├── requirements.txt +├── config/default.yaml # edit variant_feature after Step 3 +├── data/raw/ +│ ├── jm1.csv # training project +│ └── cm1.csv # evaluation project +└── scripts/ + ├── download_data.py + ├── 01_logistic_regression.py + ├── 02_gradient_boosted_trees.py + ├── 03_compare_features.py + └── 04_train_neural_network.py +``` + +## Results summary + +Typical AUC-ROC scores from this tutorial: + +| Step | Method | Split | ~AUC | +|------|--------|-------|------| +| 1 | Logistic regression | JM1 holdout | 0.727 | +| 2 | Gradient boosted trees | JM1 holdout | 0.73+ | +| 3 | Best feature (`total_Opnd`) | JM1 holdout | 0.732 | +| 4 | Neural network (`i`) | JM1 → CM1 | 0.767 | +| 4 | Neural network (`total_Opnd`) | JM1 → CM1 | **0.789** | + +Metrics are saved under `checkpoints/` after each step. + +## The six features + +From the notebook, five metrics are fixed and one is swappable: + +| # | Feature | Description | +|---|---------|-------------| +| 1 | `loc` | McCabe line count | +| 2 | `d` | Halstead difficulty | +| 3 | `locCodeAndComment` | Lines of code and comments | +| 4 | `v(g)` | Cyclomatic complexity | +| 5 | `uniq_Opnd` | Unique operands | +| 6 | *swappable* | Default `i`; tutorial picks `total_Opnd` | + +All 21 available metrics are listed in `src/data/nasa_preprocessing.py`. + +## Going further + +- Try different feature combinations in `config/default.yaml` +- Lower `training.learning_rate` if validation AUC bounces between epochs +- The neural network already uses Adam — tuning learning rate and epochs is the next lever +- If prompted by the instructor, play around with subbing out other features rather than just the 6th to see if you can further improve predictive accuracy! + +## License + +NASA PROMISE data: see [PROMISE repository](http://promise.site.uottawa.ca/SERepository) attribution guidelines. +Coveros AI for Testers training material all rights reserved +Collaborated with Professional AI Agents LLC diff --git a/defect_prediction_neural_network_tutorial/config/default.yaml b/defect_prediction_neural_network_tutorial/config/default.yaml new file mode 100644 index 0000000..98af7dc --- /dev/null +++ b/defect_prediction_neural_network_tutorial/config/default.yaml @@ -0,0 +1,42 @@ +# Shared settings for the defect prediction tutorial + +data: + train_path: data/raw/jm1.csv + val_path: data/raw/cm1.csv + label_column: defects + preprocessing: nasa + split_mode: external # used by the neural network (Step 4): train JM1, test CM1 + base_features: + - loc + - d + - locCodeAndComment + - v(g) + - uniq_Opnd + variant_feature: i # Step 3 finds the best swap; then set to total_Opnd for Step 4 + normalize: true + +model: + input_dim: 6 + hidden_dims: [7680, 7680, 3840, 1024] + dropout: 0.3 + num_classes: 1 + +training: + epochs: 20 + batch_size: 64 + learning_rate: 0.0001 + weight_decay: 0.0001 + val_split: 0.2 # used by Steps 1-3: 80/20 holdout on JM1 + seed: 42 + num_workers: 0 + checkpoint_dir: checkpoints + save_best: true + +hardware: + device: auto + +boosting: + n_estimators: 500 + learning_rate: 0.05 + max_depth: 6 + early_stopping_rounds: 50 diff --git a/defect_prediction_neural_network_tutorial/data/processed/.gitkeep b/defect_prediction_neural_network_tutorial/data/processed/.gitkeep new file mode 100644 index 0000000..ccfe662 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/data/processed/.gitkeep @@ -0,0 +1 @@ +# Placeholder for processed datasets diff --git a/defect_prediction_neural_network_tutorial/data/raw/.gitkeep b/defect_prediction_neural_network_tutorial/data/raw/.gitkeep new file mode 100644 index 0000000..fb46768 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/data/raw/.gitkeep @@ -0,0 +1 @@ +# Placeholder for raw training/validation CSV files diff --git a/defect_prediction_neural_network_tutorial/defect_prediction.ipynb b/defect_prediction_neural_network_tutorial/defect_prediction.ipynb new file mode 100644 index 0000000..0ad925d --- /dev/null +++ b/defect_prediction_neural_network_tutorial/defect_prediction.ipynb @@ -0,0 +1,200 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "246095f2-b7c0-4d92-8887-b192cf14e24b", + "metadata": {}, + "source": [ + "

Exercise 5 - Building a Defect Prediction System

" + ] + }, + { + "cell_type": "markdown", + "id": "81ec1499-54a7-4df2-9bf7-2ba4d4acf6a9", + "metadata": {}, + "source": [ + "**Introduction.** During this exercise you will be building a defect prediction system. Since creating such functionality from scratch would take more time than we have available, you will have the opportunity to select the features of the data that the model will use during both training and prediction.\n", + "\n", + "Please read the background information for each step and then execute the code block. The steps that require you to modify the code will be clearly marked.\n", + "\n", + "Thank you to prabhdeep123, who created the defect prediction system that this exercise was heavily influenced by. If you wish you can [view the original code](https://www.kaggle.com/code/prabhdeep123/software-defect-prediction)." + ] + }, + { + "cell_type": "markdown", + "id": "d2bd1d6d-f7ed-45b9-9e49-a44281098919", + "metadata": {}, + "source": [ + "**Step 1.** The first step is to import the libraries needed to implement our defect prediction approach. We then ignore warnings that we don't need to concern ourselves with." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "aca05e80-90a7-4415-aef3-e8976830cf4c", + "metadata": {}, + "outputs": [], + "source": [ + "# Import any needed libraries.\n", + "import numpy as np # linear algebra\n", + "import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n", + "from sklearn.model_selection import train_test_split\n", + "from sklearn.linear_model import LogisticRegression\n", + "from sklearn.metrics import roc_auc_score\n", + "import warnings\n", + "\n", + "# Disable all warnings\n", + "warnings.filterwarnings(\"ignore\")" + ] + }, + { + "cell_type": "markdown", + "id": "2784d6f9-8f2f-4ce6-a950-37d24a430c6a", + "metadata": {}, + "source": [ + "**Step 2.** We then load the datasets that are stored on the local filesystem. *jm1.csv* contains source code metrics and defect presence information for a NASA real-time predictive ground system written in C and will be used to train the model. *cm1.csv* contains the same information for a NASA spacecraft instrument written in C and will be used to evaluate the model.\n", + "\n", + "We then drop those rows that have missing values." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "e27e72a7-e159-41e7-8a55-d2aa463b87e6", + "metadata": {}, + "outputs": [], + "source": [ + "# Load the train and test datasets.\n", + "train_df = pd.read_csv('jm1.csv')\n", + "test_df = pd.read_csv('cm1.csv')\n", + "\n", + "# Prepare the data to be processed.\n", + "indexes = test_df.index\n", + "train_df.replace('?', pd.NA, inplace=True)\n", + "test_df.replace('?', pd.NA, inplace=True)\n", + "train_df.dropna(subset=train_df.columns[4:6], inplace=True)" + ] + }, + { + "cell_type": "markdown", + "id": "3d51dd7a-4a8c-4cd9-a082-4f87951ebecd", + "metadata": {}, + "source": [ + "**Step 3.** We then update a few columns so that they are represented as numeric values. After that we store the last column separately since information on the presence of defects will be used to evaluate the model, not to train it.\n", + "\n", + "We then select the subset of columns (features) that should be used during training and testing. **You can leave this line as-is for now, but you will later be modifying it to determine what combination of features causes the model optimizes model performance.**" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "2f656e89-9b4f-4288-abee-7c34a2a2c634", + "metadata": {}, + "outputs": [], + "source": [ + "# Make the values in the five columns preceding the final column numeric.\n", + "# Drop any rows for which such conversion fails.\n", + "train_df[train_df.columns[16:21]] = train_df[train_df.columns[16:21]].apply(pd.to_numeric, errors='coerce')\n", + "test_df[test_df.columns[16:21]] = test_df[test_df.columns[16:21]].apply(pd.to_numeric, errors='coerce')\n", + "train_df.dropna(inplace=True)\n", + "test_df.dropna(inplace=True)\n", + "train_df[train_df.columns[16:21]] = train_df[train_df.columns[16:21]].astype(int)\n", + "test_df[test_df.columns[16:21]] = test_df[test_df.columns[16:21]].astype(int)\n", + "\n", + "# Store the defects column from the training data separately.\n", + "X = train_df.drop('defects', axis = 1)\n", + "y = train_df['defects'].astype('int')\n", + "\n", + "# Select some subset of the below columns to use during training and testing.\n", + "# 1. loc : numeric % McCabe's line count of code\n", + "# 2. v(g) : numeric % McCabe \"cyclomatic complexity\"\n", + "# 3. ev(g) : numeric % McCabe \"essential complexity\"\n", + "# 4. iv(g) : numeric % McCabe \"design complexity\"\n", + "# 5. n : numeric % Halstead total operators + operands\n", + "# 6. v : numeric % Halstead \"volume\"\n", + "# 7. l : numeric % Halstead \"program length\"\n", + "# 8. d : numeric % Halstead \"difficulty\"\n", + "# 9. i : numeric % Halstead \"intelligence\"\n", + "# 10. e : numeric % Halstead \"effort\"\n", + "# 11. b : numeric % Halstead\n", + "# 12. t : numeric % Halstead's time estimator\n", + "# 13. lOCode : numeric % Halstead's line count\n", + "# 14. lOComment : numeric % Halstead's count of lines of comments\n", + "# 15. lOBlank : numeric % Halstead's count of blank lines\n", + "# 16. locCodeAndComment : numeric\n", + "# 17. uniq_Op : numeric % unique operators\n", + "# 18. uniq_Opnd : numeric % unique operands\n", + "# 19. total_Op : numeric % total operators\n", + "# 20. total_Opnd : numeric % total operands\n", + "# 21: branchCount : numeric % of the flow graph\n", + "X = X[['loc', 'd', 'locCodeAndComment', 'v(g)', 'uniq_Opnd', 'i']]\n", + "\n", + "# Prepare data to be used during training and testing.\n", + "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)" + ] + }, + { + "cell_type": "markdown", + "id": "288fca1b-3b70-4b05-a77e-10baeefebec1", + "metadata": {}, + "source": [ + "**Step 4.** The final step is to create a logistic regression model using the training data and to then evaluate it against the test data. The metric that we use to evaluate its performance is \"AUC-ROC\", which is the area under the Receiver Operating Characteristic curve. This value will be between 0 and 1, where 0 indicates a model that has not predictive power while a value of 1 means that the model perfectly predicts defects.\n", + "\n", + "**Now that you have a sense for how this defect prediction system works, you should perform some experimentation to determine what combination of features (columns) in the training data yields the highest quality predictions. You can go to Step 3, modify the list of features, and then re-run both Steps 3 and 4 to see how the AUC-ROC value changes. Alternatively, you can press 'Ctrl-F9' to re-run all of the blocks.**\n", + "\n", + "**You may wish to start by using individual features so you can see which of them have greater predictive power and then combining those that seem promising. Please be prepared to discuss what you've learned as well as the maximum AUC-ROC value you achieved with the class.**" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "50f9bac9-f809-47d2-9dc9-d1ead3c73efb", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Logistic Regression AUC-ROC: 0.7271697390458622\n" + ] + } + ], + "source": [ + "logistic_model = LogisticRegression(random_state=42)\n", + "logistic_model.fit(X_train, y_train)\n", + "logistic_predictions = logistic_model.predict_proba(X_test)[:, 1]\n", + "logistic_auc_roc = roc_auc_score(y_test, logistic_predictions)\n", + "print(f\"Logistic Regression AUC-ROC: {logistic_auc_roc}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b7814697-305d-449f-89f7-78d31e83d04a", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "training-env-python-3-12", + "language": "python", + "name": "training-env-python-3-12" + }, + "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.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/defect_prediction_neural_network_tutorial/requirements.txt b/defect_prediction_neural_network_tutorial/requirements.txt new file mode 100644 index 0000000..a8455a1 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/requirements.txt @@ -0,0 +1,5 @@ +torch>=2.2.0 +numpy>=1.26.0 +pandas>=2.1.0 +scikit-learn>=1.4.0 +pyyaml>=6.0.1 diff --git a/defect_prediction_neural_network_tutorial/scripts/01_logistic_regression.py b/defect_prediction_neural_network_tutorial/scripts/01_logistic_regression.py new file mode 100644 index 0000000..28868dc --- /dev/null +++ b/defect_prediction_neural_network_tutorial/scripts/01_logistic_regression.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +"""Step 1: Logistic regression baseline (matches defect_prediction.ipynb).""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +from sklearn.linear_model import LogisticRegression +from sklearn.metrics import roc_auc_score +from sklearn.model_selection import train_test_split + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from src.config import load_config +from src.data.dataset import resolve_feature_columns +from src.data.nasa_preprocessing import load_nasa_dataset +from src.paths import CHECKPOINT_DIR, CONFIG_PATH, resolve_path + + +def main() -> None: + config = load_config(CONFIG_PATH) + train_df = load_nasa_dataset(resolve_path(config.data.train_path)) + features = resolve_feature_columns(config) + + x = train_df[features] + y = train_df[config.data.label_column].astype(int) + + x_train, x_test, y_train, y_test = train_test_split( + x, y, test_size=config.training.val_split, random_state=config.training.seed + ) + + model = LogisticRegression(random_state=config.training.seed, max_iter=1000) + model.fit(x_train, y_train) + auc = roc_auc_score(y_test, model.predict_proba(x_test)[:, 1]) + + print("Step 1: Logistic Regression") + print(f" Features: {features}") + print(f" AUC-ROC: {auc:.4f} (~0.727 with default features)") + + out = CHECKPOINT_DIR / "01_logistic_regression.json" + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w") as f: + json.dump({"features": features, "auc_roc": auc}, f, indent=2) + print(f" Saved: {out.relative_to(ROOT)}") + + +if __name__ == "__main__": + main() diff --git a/defect_prediction_neural_network_tutorial/scripts/02_gradient_boosted_trees.py b/defect_prediction_neural_network_tutorial/scripts/02_gradient_boosted_trees.py new file mode 100644 index 0000000..08ff40c --- /dev/null +++ b/defect_prediction_neural_network_tutorial/scripts/02_gradient_boosted_trees.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +"""Step 2: Gradient boosted trees baseline (sklearn, no extra installs).""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +from sklearn.ensemble import HistGradientBoostingClassifier +from sklearn.metrics import roc_auc_score +from sklearn.model_selection import train_test_split + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from src.config import load_config +from src.data.dataset import resolve_feature_columns +from src.data.nasa_preprocessing import load_nasa_dataset +from src.paths import CHECKPOINT_DIR, CONFIG_PATH, resolve_path + + +def main() -> None: + config = load_config(CONFIG_PATH) + train_df = load_nasa_dataset(resolve_path(config.data.train_path)) + features = resolve_feature_columns(config) + + x = train_df[features].to_numpy() + y = train_df[config.data.label_column].to_numpy(dtype=int) + + x_train, x_test, y_train, y_test = train_test_split( + x, y, test_size=config.training.val_split, random_state=config.training.seed + ) + + pos = float(np.sum(y_train == 1)) + neg = float(np.sum(y_train == 0)) + sample_weight = np.where(y_train == 1, neg / pos, 1.0) if pos else None + + model = HistGradientBoostingClassifier( + max_iter=config.boosting.n_estimators, + learning_rate=config.boosting.learning_rate, + max_depth=config.boosting.max_depth, + random_state=config.training.seed, + early_stopping=True, + n_iter_no_change=config.boosting.early_stopping_rounds, + ) + model.fit(x_train, y_train, sample_weight=sample_weight) + auc = roc_auc_score(y_test, model.predict_proba(x_test)[:, 1]) + + print("Step 2: Gradient Boosted Trees") + print(f" Features: {features}") + print(f" AUC-ROC: {auc:.4f}") + + out = CHECKPOINT_DIR / "02_gradient_boosted_trees.json" + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w") as f: + json.dump({"features": features, "auc_roc": auc}, f, indent=2) + print(f" Saved: {out.relative_to(ROOT)}") + + +if __name__ == "__main__": + main() diff --git a/defect_prediction_neural_network_tutorial/scripts/03_compare_features.py b/defect_prediction_neural_network_tutorial/scripts/03_compare_features.py new file mode 100644 index 0000000..9b549f7 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/scripts/03_compare_features.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Step 3: Try each candidate for the 6th feature and pick the best.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +from sklearn.linear_model import LogisticRegression +from sklearn.metrics import roc_auc_score +from sklearn.model_selection import train_test_split + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from src.config import load_config +from src.data.nasa_preprocessing import ( + NASA_BASE_FEATURES, + build_feature_columns, + load_nasa_dataset, + variant_feature_options, +) +from src.paths import CHECKPOINT_DIR, CONFIG_PATH, resolve_path + + +def main() -> None: + config = load_config(CONFIG_PATH) + train_df = load_nasa_dataset(resolve_path(config.data.train_path)) + base = list(config.data.base_features or NASA_BASE_FEATURES) + + print("Step 3: Compare 6th Feature Candidates") + print(f" Fixed features: {base}\n") + print(f"{'Rank':<5} {'Feature':<16} {'AUC-ROC':>8}") + print("-" * 32) + + results = [] + for variant in variant_feature_options(base): + features = build_feature_columns(base, variant) + x = train_df[features] + y = train_df[config.data.label_column].astype(int) + + x_train, x_test, y_train, y_test = train_test_split( + x, y, test_size=config.training.val_split, random_state=config.training.seed + ) + model = LogisticRegression(random_state=config.training.seed, max_iter=1000) + model.fit(x_train, y_train) + auc = roc_auc_score(y_test, model.predict_proba(x_test)[:, 1]) + results.append({"variant_feature": variant, "features": features, "auc_roc": auc}) + + results.sort(key=lambda r: r["auc_roc"], reverse=True) + for rank, row in enumerate(results, start=1): + mark = " <-- best" if rank == 1 else "" + print(f"{rank:<5} {row['variant_feature']:<16} {row['auc_roc']:>8.4f}{mark}") + + best = results[0] + print(f"\n Update config/default.yaml:") + print(f" variant_feature: {best['variant_feature']}") + print(f" Then run Step 4.") + + out = CHECKPOINT_DIR / "03_feature_comparison.json" + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w") as f: + json.dump({"base_features": base, "results": results, "best": best}, f, indent=2) + print(f" Saved: {out.relative_to(ROOT)}") + + +if __name__ == "__main__": + main() diff --git a/defect_prediction_neural_network_tutorial/scripts/04_train_neural_network.py b/defect_prediction_neural_network_tutorial/scripts/04_train_neural_network.py new file mode 100644 index 0000000..7913193 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/scripts/04_train_neural_network.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +"""Step 4: Train the neural network on JM1 and evaluate on CM1.""" + +from __future__ import annotations + +import pickle +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from src.config import load_config +from src.data.dataset import build_dataloader, load_splits +from src.model.defect_predictor import DefectPredictor, count_parameters +from src.paths import CONFIG_PATH, resolve_path +from src.training.trainer import Trainer + + +def main() -> None: + config = load_config(CONFIG_PATH) + train_ds, val_ds, scaler, features = load_splits(config) + + config.model.input_dim = len(features) + model = DefectPredictor(config.model) + + print("Step 4: Neural Network") + print(f" Features: {features}") + print(f" Parameters: {count_parameters(model):,}") + print(f" Train: {config.data.train_path}") + print(f" Evaluate: {config.data.val_path}") + print() + + trainer = Trainer( + model, + config, + build_dataloader(train_ds, config, True), + build_dataloader(val_ds, config, False), + ) + history = trainer.train() + + best = max(history, key=lambda e: e.get("val_roc_auc", 0)) + print(f"\n Best val AUC: {best['val_roc_auc']:.4f} (epoch {best['epoch']})") + + if scaler is not None: + path = resolve_path(config.training.checkpoint_dir) / "scaler.pkl" + with open(path, "wb") as f: + pickle.dump(scaler, f) + + +if __name__ == "__main__": + main() diff --git a/defect_prediction_neural_network_tutorial/scripts/download_data.py b/defect_prediction_neural_network_tutorial/scripts/download_data.py new file mode 100644 index 0000000..0162c5e --- /dev/null +++ b/defect_prediction_neural_network_tutorial/scripts/download_data.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 +"""Download NASA PROMISE JM1 and CM1 datasets.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from urllib.request import urlretrieve + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from src.paths import CM1_PATH, DATA_DIR, JM1_PATH + +DATASETS = { + JM1_PATH: ( + "https://raw.githubusercontent.com/ApoorvaKrisna/" + "NASA-promise-dataset-repository/main/jm1.csv" + ), + CM1_PATH: ( + "https://raw.githubusercontent.com/ApoorvaKrisna/" + "NASA-promise-dataset-repository/main/cm1.csv" + ), +} + + +def main() -> None: + DATA_DIR.mkdir(parents=True, exist_ok=True) + + for dest, url in DATASETS.items(): + print(f"Downloading {dest.name} ...") + urlretrieve(url, dest) + print(f" -> {dest.relative_to(ROOT)}") + + +if __name__ == "__main__": + main() diff --git a/defect_prediction_neural_network_tutorial/src/__init__.py b/defect_prediction_neural_network_tutorial/src/__init__.py new file mode 100644 index 0000000..30062bd --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/__init__.py @@ -0,0 +1 @@ +"""Defect prediction neural network training package.""" diff --git a/defect_prediction_neural_network_tutorial/src/config.py b/defect_prediction_neural_network_tutorial/src/config.py new file mode 100644 index 0000000..660fb6c --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/config.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import yaml + + +@dataclass +class DataConfig: + train_path: str = "data/raw/jm1.csv" + val_path: str | None = "data/raw/cm1.csv" + label_column: str = "defects" + feature_columns: list[str] | None = None + base_features: list[str] | None = None + variant_feature: str = "i" + normalize: bool = True + preprocessing: str = "nasa" # nasa | none + split_mode: str = "external" # external | random + + +@dataclass +class ModelConfig: + input_dim: int = 1536 + hidden_dims: list[int] = field(default_factory=lambda: [7168, 7168, 3584, 1024]) + dropout: float = 0.2 + num_classes: int = 1 + + +@dataclass +class TrainingConfig: + epochs: int = 20 + batch_size: int = 64 + learning_rate: float = 1e-4 + weight_decay: float = 1e-5 + val_split: float = 0.2 + seed: int = 42 + num_workers: int = 0 + checkpoint_dir: str = "checkpoints" + save_best: bool = True + + +@dataclass +class HardwareConfig: + device: str = "auto" + + +@dataclass +class BoostingConfig: + n_estimators: int = 500 + learning_rate: float = 0.05 + max_depth: int = 6 + early_stopping_rounds: int = 50 + + +@dataclass +class Config: + data: DataConfig = field(default_factory=DataConfig) + model: ModelConfig = field(default_factory=ModelConfig) + training: TrainingConfig = field(default_factory=TrainingConfig) + hardware: HardwareConfig = field(default_factory=HardwareConfig) + boosting: BoostingConfig = field(default_factory=BoostingConfig) + + +def _merge_dict(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: + merged = dict(base) + for key, value in override.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = _merge_dict(merged[key], value) + else: + merged[key] = value + return merged + + +def _section(data: dict[str, Any], cls: type): + return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__}) + + +def load_config(path: str | Path) -> Config: + with open(path) as f: + raw = yaml.safe_load(f) or {} + + return Config( + data=_section(raw.get("data", {}), DataConfig), + model=_section(raw.get("model", {}), ModelConfig), + training=_section(raw.get("training", {}), TrainingConfig), + hardware=_section(raw.get("hardware", {}), HardwareConfig), + boosting=_section(raw.get("boosting", {}), BoostingConfig), + ) diff --git a/defect_prediction_neural_network_tutorial/src/data/__init__.py b/defect_prediction_neural_network_tutorial/src/data/__init__.py new file mode 100644 index 0000000..0969b4f --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/data/__init__.py @@ -0,0 +1,24 @@ +from .dataset import DefectDataset, build_dataloader, load_splits, resolve_feature_columns +from .nasa_preprocessing import ( + NASA_ALL_FEATURES, + NASA_BASE_FEATURES, + NASA_DEFAULT_FEATURES, + NASA_DEFAULT_VARIANT_FEATURE, + build_feature_columns, + load_nasa_dataset, + variant_feature_options, +) + +__all__ = [ + "DefectDataset", + "NASA_ALL_FEATURES", + "NASA_BASE_FEATURES", + "NASA_DEFAULT_FEATURES", + "NASA_DEFAULT_VARIANT_FEATURE", + "build_dataloader", + "build_feature_columns", + "load_nasa_dataset", + "load_splits", + "resolve_feature_columns", + "variant_feature_options", +] diff --git a/defect_prediction_neural_network_tutorial/src/data/dataset.py b/defect_prediction_neural_network_tutorial/src/data/dataset.py new file mode 100644 index 0000000..38d40b2 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/data/dataset.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from sklearn.model_selection import train_test_split +from sklearn.preprocessing import StandardScaler +from torch.utils.data import DataLoader, Dataset + +from src.config import Config +from src.paths import resolve_path +from src.data.nasa_preprocessing import ( + NASA_DEFAULT_FEATURES, + NASA_DEFAULT_VARIANT_FEATURE, + build_feature_columns, + load_nasa_dataset, + variant_feature_options, +) + + +class DefectDataset(Dataset): + def __init__(self, features: np.ndarray, labels: np.ndarray): + self.features = torch.as_tensor(features, dtype=torch.float32) + self.labels = torch.as_tensor(labels, dtype=torch.float32) + + def __len__(self) -> int: + return len(self.features) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: + return self.features[idx], self.labels[idx] + + +def _read_csv(path: Path, config: Config) -> pd.DataFrame: + if not path.exists(): + raise FileNotFoundError(f"Dataset not found: {path}") + if config.data.preprocessing == "nasa": + return load_nasa_dataset(str(path)) + return pd.read_csv(path) + + +def _extract_xy( + df: pd.DataFrame, + label_column: str, + feature_columns: list[str] | None, +) -> tuple[np.ndarray, np.ndarray]: + if label_column not in df.columns: + raise ValueError(f"Label column '{label_column}' not found in dataset") + + if feature_columns is None: + feature_columns = [c for c in df.columns if c != label_column] + + missing = [c for c in feature_columns if c not in df.columns] + if missing: + raise ValueError(f"Missing feature columns: {missing}") + + x = df[feature_columns].to_numpy(dtype=np.float32) + y = df[label_column].to_numpy(dtype=np.float32) + + if y.ndim == 1: + y = y.reshape(-1, 1) + + return x, y + + +def resolve_feature_columns(config: Config) -> list[str]: + if config.data.feature_columns: + return list(config.data.feature_columns) + if config.data.preprocessing == "nasa": + return build_feature_columns( + base_features=config.data.base_features, + variant_feature=config.data.variant_feature, + ) + return [] + + +def load_splits( + config: Config, +) -> tuple[DefectDataset, DefectDataset, StandardScaler | None, list[str]]: + feature_columns = resolve_feature_columns(config) + train_df = _read_csv(resolve_path(config.data.train_path), config) + val_path = resolve_path(config.data.val_path) if config.data.val_path else None + + if config.data.split_mode == "external" and val_path and val_path.exists(): + val_df = _read_csv(val_path, config) + x_train, y_train = _extract_xy(train_df, config.data.label_column, feature_columns) + x_val, y_val = _extract_xy(val_df, config.data.label_column, feature_columns) + else: + x_all, y_all = _extract_xy(train_df, config.data.label_column, feature_columns) + x_train, x_val, y_train, y_val = train_test_split( + x_all, + y_all, + test_size=config.training.val_split, + random_state=config.training.seed, + stratify=y_all if len(np.unique(y_all)) > 1 else None, + ) + + if not feature_columns: + feature_columns = [ + c for c in train_df.columns if c != config.data.label_column + ] + + scaler: StandardScaler | None = None + if config.data.normalize: + scaler = StandardScaler() + x_train = scaler.fit_transform(x_train).astype(np.float32) + x_val = scaler.transform(x_val).astype(np.float32) + + return ( + DefectDataset(x_train, y_train), + DefectDataset(x_val, y_val), + scaler, + feature_columns, + ) + + +def build_dataloader(dataset: DefectDataset, config: Config, shuffle: bool) -> DataLoader: + return DataLoader( + dataset, + batch_size=config.training.batch_size, + shuffle=shuffle, + num_workers=config.training.num_workers, + pin_memory=False, + ) diff --git a/defect_prediction_neural_network_tutorial/src/data/nasa_preprocessing.py b/defect_prediction_neural_network_tutorial/src/data/nasa_preprocessing.py new file mode 100644 index 0000000..dd3600b --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/data/nasa_preprocessing.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +import pandas as pd + +# All software metrics available in the NASA PROMISE JM1/CM1 datasets. +NASA_ALL_FEATURES: list[str] = [ + "loc", + "v(g)", + "ev(g)", + "iv(g)", + "n", + "v", + "l", + "d", + "i", + "e", + "b", + "t", + "lOCode", + "lOComment", + "lOBlank", + "locCodeAndComment", + "uniq_Op", + "uniq_Opnd", + "total_Op", + "total_Opnd", + "branchCount", +] + +# Default feature subset from defect_prediction.ipynb (Step 3). +NASA_BASE_FEATURES: list[str] = [ + "loc", + "d", + "locCodeAndComment", + "v(g)", + "uniq_Opnd", +] + +NASA_DEFAULT_VARIANT_FEATURE: str = "i" + +NASA_DEFAULT_FEATURES: list[str] = NASA_BASE_FEATURES + [NASA_DEFAULT_VARIANT_FEATURE] + + +def build_feature_columns( + base_features: list[str] | None = None, + variant_feature: str = NASA_DEFAULT_VARIANT_FEATURE, +) -> list[str]: + """Build the 6-feature set: 5 fixed metrics + one swappable metric.""" + base = list(base_features or NASA_BASE_FEATURES) + if variant_feature in base: + raise ValueError( + f"variant_feature '{variant_feature}' is already in base_features: {base}" + ) + if variant_feature not in NASA_ALL_FEATURES: + raise ValueError( + f"Unknown variant_feature '{variant_feature}'. " + f"Choose from: {NASA_ALL_FEATURES}" + ) + return base + [variant_feature] + + +def variant_feature_options( + base_features: list[str] | None = None, +) -> list[str]: + """Metrics available for the 6th feature slot (excludes base features).""" + base = set(base_features or NASA_BASE_FEATURES) + return [name for name in NASA_ALL_FEATURES if name not in base] + + +def preprocess_nasa_dataframe(df: pd.DataFrame) -> pd.DataFrame: + """Apply the cleaning steps from defect_prediction.ipynb.""" + cleaned = df.copy() + cleaned.replace("?", pd.NA, inplace=True) + cleaned.dropna(subset=cleaned.columns[4:6], inplace=True) + + numeric_cols = cleaned.columns[16:21] + cleaned[numeric_cols] = cleaned[numeric_cols].apply(pd.to_numeric, errors="coerce") + cleaned.dropna(inplace=True) + cleaned[numeric_cols] = cleaned[numeric_cols].astype(int) + + if cleaned["defects"].dtype == object: + cleaned["defects"] = ( + cleaned["defects"] + .astype(str) + .str.strip() + .str.lower() + .map({"false": 0, "true": 1}) + ) + else: + cleaned["defects"] = cleaned["defects"].astype(int) + + cleaned.dropna(subset=["defects"], inplace=True) + cleaned["defects"] = cleaned["defects"].astype(int) + return cleaned.reset_index(drop=True) + + +def load_nasa_dataset(path: str) -> pd.DataFrame: + return preprocess_nasa_dataframe(pd.read_csv(path)) diff --git a/defect_prediction_neural_network_tutorial/src/model/__init__.py b/defect_prediction_neural_network_tutorial/src/model/__init__.py new file mode 100644 index 0000000..5534348 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/model/__init__.py @@ -0,0 +1,3 @@ +from .defect_predictor import DefectPredictor, count_parameters + +__all__ = ["DefectPredictor", "count_parameters"] diff --git a/defect_prediction_neural_network_tutorial/src/model/defect_predictor.py b/defect_prediction_neural_network_tutorial/src/model/defect_predictor.py new file mode 100644 index 0000000..01d3a65 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/model/defect_predictor.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import torch +import torch.nn as nn + +from src.config import ModelConfig + + +class DefectPredictor(nn.Module): + """MLP for binary defect prediction.""" + + def __init__(self, config: ModelConfig): + super().__init__() + self.config = config + + layers: list[nn.Module] = [] + in_dim = config.input_dim + + for hidden_dim in config.hidden_dims: + layers.extend( + [ + nn.Linear(in_dim, hidden_dim), + nn.BatchNorm1d(hidden_dim), + nn.GELU(), + nn.Dropout(config.dropout), + ] + ) + in_dim = hidden_dim + + layers.append(nn.Linear(in_dim, config.num_classes)) + self.network = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.network(x) + + +def count_parameters(model: nn.Module) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def build_model(config: ModelConfig) -> DefectPredictor: + return DefectPredictor(config) diff --git a/defect_prediction_neural_network_tutorial/src/paths.py b/defect_prediction_neural_network_tutorial/src/paths.py new file mode 100644 index 0000000..c1d96f7 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/paths.py @@ -0,0 +1,21 @@ +"""Project-root paths that work no matter where a script is run from.""" + +from __future__ import annotations + +from pathlib import Path + +# Repo root: parent of src/ +ROOT = Path(__file__).resolve().parents[1] + +DATA_DIR = ROOT / "data" / "raw" +CONFIG_PATH = ROOT / "config" / "default.yaml" +CHECKPOINT_DIR = ROOT / "checkpoints" + +JM1_PATH = DATA_DIR / "jm1.csv" +CM1_PATH = DATA_DIR / "cm1.csv" + + +def resolve_path(path: str | Path) -> Path: + """Turn a config-relative path into an absolute path under ROOT.""" + p = Path(path) + return p if p.is_absolute() else ROOT / p diff --git a/defect_prediction_neural_network_tutorial/src/training/__init__.py b/defect_prediction_neural_network_tutorial/src/training/__init__.py new file mode 100644 index 0000000..81618ca --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/training/__init__.py @@ -0,0 +1,3 @@ +from .trainer import Trainer + +__all__ = ["Trainer"] diff --git a/defect_prediction_neural_network_tutorial/src/training/trainer.py b/defect_prediction_neural_network_tutorial/src/training/trainer.py new file mode 100644 index 0000000..b68dbb8 --- /dev/null +++ b/defect_prediction_neural_network_tutorial/src/training/trainer.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import json +from dataclasses import asdict +from pathlib import Path + +import torch +import torch.nn as nn +from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score +from torch.utils.data import DataLoader +from src.config import Config +from src.paths import resolve_path +from src.model.defect_predictor import DefectPredictor, count_parameters + + +def resolve_device(preference: str) -> torch.device: + if preference == "auto": + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + return torch.device(preference) + + +class Trainer: + def __init__( + self, + model: DefectPredictor, + config: Config, + train_loader: DataLoader, + val_loader: DataLoader, + ): + self.model = model + self.config = config + self.train_loader = train_loader + self.val_loader = val_loader + self.device = resolve_device(config.hardware.device) + self.model.to(self.device) + + self.criterion = nn.BCEWithLogitsLoss() + self.optimizer = torch.optim.AdamW( + self.model.parameters(), + lr=config.training.learning_rate, + weight_decay=config.training.weight_decay, + ) + + self.checkpoint_dir = resolve_path(config.training.checkpoint_dir) + self.checkpoint_dir.mkdir(parents=True, exist_ok=True) + self.history: list[dict[str, float]] = [] + self.best_val_loss = float("inf") + + def _run_epoch(self, loader: DataLoader, train: bool) -> tuple[float, dict[str, float]]: + self.model.train(train) + total_loss = 0.0 + all_probs: list[float] = [] + all_labels: list[float] = [] + + context = torch.enable_grad() if train else torch.no_grad() + with context: + for features, labels in loader: + features = features.to(self.device) + labels = labels.to(self.device) + + logits = self.model(features) + loss = self.criterion(logits, labels) + + if train: + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + + total_loss += loss.item() * features.size(0) + probs = torch.sigmoid(logits).detach().cpu().numpy().ravel() + all_probs.extend(probs.tolist()) + all_labels.extend(labels.detach().cpu().numpy().ravel().tolist()) + + avg_loss = total_loss / len(loader.dataset) + preds = [1 if p >= 0.5 else 0 for p in all_probs] + metrics = { + "accuracy": float(accuracy_score(all_labels, preds)), + "precision": float(precision_score(all_labels, preds, zero_division=0)), + "recall": float(recall_score(all_labels, preds, zero_division=0)), + "f1": float(f1_score(all_labels, preds, zero_division=0)), + } + + if len(set(all_labels)) > 1: + metrics["roc_auc"] = float(roc_auc_score(all_labels, all_probs)) + else: + metrics["roc_auc"] = float("nan") + + return avg_loss, metrics + + def train(self) -> list[dict[str, float]]: + param_count = count_parameters(self.model) + print(f"Model parameters: {param_count:,}") + print(f"Training on {self.device} for {self.config.training.epochs} epochs") + + for epoch in range(1, self.config.training.epochs + 1): + train_loss, train_metrics = self._run_epoch(self.train_loader, train=True) + val_loss, val_metrics = self._run_epoch(self.val_loader, train=False) + + record = { + "epoch": epoch, + "train_loss": train_loss, + "val_loss": val_loss, + **{f"train_{k}": v for k, v in train_metrics.items()}, + **{f"val_{k}": v for k, v in val_metrics.items()}, + } + self.history.append(record) + + print( + f"Epoch {epoch:02d}/{self.config.training.epochs} | " + f"train_loss={train_loss:.4f} val_loss={val_loss:.4f} | " + f"val_f1={val_metrics['f1']:.4f} val_auc={val_metrics['roc_auc']:.4f}" + ) + + if self.config.training.save_best and val_loss < self.best_val_loss: + self.best_val_loss = val_loss + self._save_checkpoint("best.pt", record) + + self._save_checkpoint("last.pt", self.history[-1]) + self._save_history() + return self.history + + def _save_checkpoint(self, filename: str, metrics: dict) -> None: + path = self.checkpoint_dir / filename + torch.save( + { + "model_state_dict": self.model.state_dict(), + "config": asdict(self.config), + "metrics": metrics, + }, + path, + ) + + def _save_history(self) -> None: + history_path = self.checkpoint_dir / "history.json" + with open(history_path, "w") as f: + json.dump(self.history, f, indent=2)