{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":53666,"databundleVersionId":6589269},{"sourceType":"datasetVersion","sourceId":16052560,"datasetId":10227749,"databundleVersionId":17019901},{"sourceType":"modelInstanceVersion","sourceId":850269,"databundleVersionId":16946521,"modelInstanceId":646411,"modelId":658355}],"dockerImageVersionId":31328,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"colab":{"provenance":[]}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Sleep and Awake Detection with MLP using Acceleration Signals 🛌\nThis notebook trains a neural network to detect sleep and awake events from accelerometer time-series data (Child Mind Institute's Detect Sleep States competition) in an efficient way. Each epoch of training and validation on around 100 million records takes 15 seconds thanks to GPU acceleration.","metadata":{"id":"shfWrOuUhOJR"}},{"cell_type":"code","source":"!pip install -q torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu118","metadata":{"trusted":true,"id":"XVT4LAQShOJR","outputId":"8e82aca8-782f-4b47-bee0-8240ab5f6983","execution":{"iopub.status.busy":"2026-05-04T08:29:54.187234Z","iopub.execute_input":"2026-05-04T08:29:54.187689Z","iopub.status.idle":"2026-05-04T08:29:57.927981Z","shell.execute_reply.started":"2026-05-04T08:29:54.187657Z","shell.execute_reply":"2026-05-04T08:29:57.927115Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom itertools import groupby\nimport gc\nimport matplotlib.pyplot as plt\nimport random\nfrom tqdm import tqdm\nimport pickle\nimport os\nfrom glob import glob\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score\nimport time\nimport optuna\nimport matplotlib.pyplot as plt\nimport shap\n\nimport torch\nfrom torch import nn\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"id":"KbX8r0sihOJS","execution":{"iopub.status.busy":"2026-05-04T08:29:57.929583Z","iopub.execute_input":"2026-05-04T08:29:57.930166Z","iopub.status.idle":"2026-05-04T08:30:01.600209Z","shell.execute_reply.started":"2026-05-04T08:29:57.930130Z","shell.execute_reply":"2026-05-04T08:30:01.599566Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"PyTorch version:\", torch.__version__)\nprint(\"CUDA available:\", torch.cuda.is_available())\nprint(\"CUDA version:\", torch.version.cuda)\nprint(\"Supported arches:\", torch.cuda.get_arch_list())\n\nif torch.cuda.is_available():\n    print(\"GPU name:\", torch.cuda.get_device_name(0))\n    print(\"Compute capability:\", torch.cuda.get_device_capability(0))","metadata":{"trusted":true,"id":"ZIxtA-t6hOJS","outputId":"f5f067f3-7313-4f08-b12f-ea344e43ca71","execution":{"iopub.status.busy":"2026-05-04T08:30:01.601154Z","iopub.execute_input":"2026-05-04T08:30:01.601769Z","iopub.status.idle":"2026-05-04T08:30:01.631073Z","shell.execute_reply.started":"2026-05-04T08:30:01.601725Z","shell.execute_reply":"2026-05-04T08:30:01.630378Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CLASSES = {\"onset\": 1, \"wakeup\": 0}\nwindow_widths = [1, 5, 30, 60, 120] # mins\nSEED = 42\nlearning_rate = 0.001\nweight_decay = 0.0001\nepochs = 20\nbatch_size = 100_000\nhidden_features = [32, 32, 32, 32]\nin_features = 77\nout_features = 1\ndebug = False\nuse_amp = True # Automatic mixed precision\nprocessed_data_dir = \"processed_CMI_data_v2.2\"\nprocessed_data_exists = os.path.exists(f\"/kaggle/input/cmi-processed-data/{processed_data_dir}\")\nvisualize = False\nnorm_type = \"z-score\"","metadata":{"trusted":true,"id":"R48UbXlohOJT","execution":{"iopub.status.busy":"2026-05-04T08:30:01.632741Z","iopub.execute_input":"2026-05-04T08:30:01.633088Z","iopub.status.idle":"2026-05-04T08:30:01.643339Z","shell.execute_reply.started":"2026-05-04T08:30:01.633063Z","shell.execute_reply":"2026-05-04T08:30:01.642662Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1. Load The Raw Dataset\nThe raw time-series data and event annotations are loaded from parquet files.","metadata":{"id":"0Emaln8shOJT"}},{"cell_type":"code","source":"if not processed_data_exists:\n    all_train_events = pd.read_csv(\"/kaggle/input/competitions/child-mind-institute-detect-sleep-states/train_events.csv\")\n\n    if not os.path.exists(f\"/kaggle/input/datasets/hadighahroudi/cmi-processed-data/train_series_with_hour.parquet\"):\n        all_train_series = pd.read_parquet(\"/kaggle/input/competitions/child-mind-institute-detect-sleep-states/train_series.parquet\")\n        all_train_series[\"hour\"] = pd.to_datetime(all_train_series[\"timestamp\"], format=\"ISO8601\",utc=True, yearfirst=True, errors=\"coerce\",cache=True).dt.tz_convert('us/eastern').dt.hour\n        all_train_series.to_parquet(\"train_series_with_hour.parquet\")\n    else:\n        all_train_series = pd.read_parquet(\"/kaggle/input/datasets/hadighahroudi/cmi-processed-data/train_series_with_hour.parquet\")","metadata":{"trusted":true,"id":"C3YnuyOrhOJT","execution":{"iopub.status.busy":"2026-05-04T08:30:01.644371Z","iopub.execute_input":"2026-05-04T08:30:01.644646Z","iopub.status.idle":"2026-05-04T08:30:01.650991Z","shell.execute_reply.started":"2026-05-04T08:30:01.644621Z","shell.execute_reply":"2026-05-04T08:30:01.650420Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Clean and Merge Time-Series Data with Events\n\nThis step removes missing event labels and formats the event data properly. The accelerometer data is merged with the cleaned event labels based on series_id and step.","metadata":{"id":"NiD2VLRPhOJT"}},{"cell_type":"code","source":"if not processed_data_exists:\n    all_train_events_no_nan = all_train_events.dropna().copy()\n    all_train_events_no_nan[\"step\"] = all_train_events_no_nan[\"step\"].astype(\"int\")\n    all_train_events_no_nan[\"event\"] = all_train_events_no_nan[\"event\"].replace(CLASSES).astype(\"int\")\n\n    train_set = pd.merge(all_train_series, all_train_events_no_nan[[\"series_id\", \"night\", \"event\", \"step\"]], on = [\"series_id\", \"step\"], how = \"left\").sort_values(by = [\"series_id\", \"step\"])\n\n    series_ids = train_set[\"series_id\"].unique()","metadata":{"trusted":true,"id":"VfO4nJ4jhOJT","execution":{"iopub.status.busy":"2026-05-04T08:30:01.652047Z","iopub.execute_input":"2026-05-04T08:30:01.652406Z","iopub.status.idle":"2026-05-04T08:30:01.663180Z","shell.execute_reply.started":"2026-05-04T08:30:01.652381Z","shell.execute_reply":"2026-05-04T08:30:01.662534Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not processed_data_exists:\n    del all_train_series, all_train_events, all_train_events_no_nan\n    gc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:30:01.664050Z","iopub.execute_input":"2026-05-04T08:30:01.664525Z","iopub.status.idle":"2026-05-04T08:30:01.673931Z","shell.execute_reply.started":"2026-05-04T08:30:01.664464Z","shell.execute_reply":"2026-05-04T08:30:01.673127Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2.1. Data Visualizations","metadata":{}},{"cell_type":"code","source":"if visualize:\n    i, j = 10000, 200000\n    slice_df = train_set[i:j]\n    plt.figure(figsize = (12, 3))\n    plt.plot(slice_df[\"step\"], slice_df[\"enmo\"])\n    ymin, ymax = plt.ylim()\n    plt.vlines(slice_df[slice_df[\"event\"] == CLASSES[\"onset\"]][\"step\"], ymin = ymin, ymax = ymax, colors= \"red\", linestyles=\"--\", label = \"onset\")\n    plt.vlines(slice_df[slice_df[\"event\"] == CLASSES[\"wakeup\"]][\"step\"], ymin = ymin, ymax = ymax, colors= \"orange\", linestyles=\"--\", label = \"wakeup\")\n    plt.title(\"Sample ENMO signal with sleep onset and wakeup events\")\n    plt.legend()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:30:01.674872Z","iopub.execute_input":"2026-05-04T08:30:01.675205Z","iopub.status.idle":"2026-05-04T08:30:01.684218Z","shell.execute_reply.started":"2026-05-04T08:30:01.675169Z","shell.execute_reply":"2026-05-04T08:30:01.683409Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if visualize:\n    sleep_durations = []\n    \n    for (sid, night), group in train_set.groupby([\"series_id\", \"night\"]):\n        group = group.sort_values(\"step\")\n    \n        onset_steps = group[group[\"event\"] == CLASSES[\"onset\"]][\"step\"].values\n        wake_steps  = group[group[\"event\"] == CLASSES[\"wakeup\"]][\"step\"].values\n    \n        i, j = 0, 0\n    \n        while i < len(onset_steps) and j < len(wake_steps):\n            onset = onset_steps[i]\n    \n            # find first wakeup after onset\n            while j < len(wake_steps) and wake_steps[j] < onset:\n                j += 1\n    \n            if j < len(wake_steps):\n                wake = wake_steps[j]\n                duration = wake - onset\n    \n                sleep_durations.append(duration)\n    \n                i += 1\n                j += 1\n            else:\n                break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:30:01.685124Z","iopub.execute_input":"2026-05-04T08:30:01.685427Z","iopub.status.idle":"2026-05-04T08:30:01.698202Z","shell.execute_reply.started":"2026-05-04T08:30:01.685392Z","shell.execute_reply":"2026-05-04T08:30:01.697637Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if visualize:\n    df_sleep = pd.DataFrame(sleep_periods)\n    # convert steps → hours (30s per step)\n    df_sleep[\"duration_hours\"] = df_sleep[\"duration_steps\"] * 5 / 3600\n    # Per-night total sleep\n    # night_sleep = df_sleep.groupby([\"series_id\", \"night\"])[\"duration_hours\"].sum().reset_index()\n\n    plt.figure(figsize=(10, 5))\n    \n    plt.hist(df_sleep[\"duration_hours\"], bins=40, label = \"Hours slept per episode\" )\n    # plt.hist(night_sleep[\"duration_hours\"], bins=40, label = \"Hours slept per night\")\n    \n    plt.title(\"Distribution of Sleep Episode Durations\")\n    plt.ylabel(\"Frequency\")\n    \n    plt.grid(alpha=0.3)\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:30:01.700587Z","iopub.execute_input":"2026-05-04T08:30:01.700870Z","iopub.status.idle":"2026-05-04T08:30:01.708938Z","shell.execute_reply.started":"2026-05-04T08:30:01.700848Z","shell.execute_reply":"2026-05-04T08:30:01.708216Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if visualize:\n    example_series = random.choice(df_sleep[\"series_id\"])\n    subset = df_sleep[df_sleep[\"series_id\"] == example_series]\n    \n    plt.figure(figsize=(10, 4))\n    \n    plt.plot(subset[\"night\"], subset[\"duration_hours\"], marker=\"o\")\n    \n    plt.title(f\"Sleep Episode Duration Across Nights\\nSeries {example_series}\")\n    plt.xlabel(\"Night\")\n    plt.ylabel(\"Hours\")\n    \n    plt.grid(alpha=0.3)\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:30:01.709841Z","iopub.execute_input":"2026-05-04T08:30:01.710225Z","iopub.status.idle":"2026-05-04T08:30:01.721922Z","shell.execute_reply.started":"2026-05-04T08:30:01.710201Z","shell.execute_reply":"2026-05-04T08:30:01.721170Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Feature extraction\nCompute window-based features for each time series. It applies rolling statistics across multiple window sizes and stacks them into feature matrices.","metadata":{"id":"NBypPbQzhOJT"}},{"cell_type":"code","source":"feature_names = [\"hour_sin\", \"hour_cos\"]\n\nfor mins in window_widths:\n    suffix = f\"{mins}min\"\n\n    feature_names.extend([\n        # enmo\n        f\"enmo_mean_abs_{suffix}\",\n        f\"enmo_max_abs_{suffix}\",\n        f\"enmo_std_abs_{suffix}\",\n\n        # enmo diff\n        f\"enmo_diff_abs_mean_{suffix}\",\n        f\"enmo_diff_abs_max_{suffix}\",\n        f\"enmo_diff_abs_std_{suffix}\",\n\n        # anglez\n        f\"anglez_mean_abs_{suffix}\",\n        f\"anglez_max_abs_{suffix}\",\n        f\"anglez_std_abs_{suffix}\",\n\n        # anglez diff\n        f\"anglez_diff_abs_mean_{suffix}\",\n        f\"anglez_diff_abs_max_{suffix}\",\n        f\"anglez_diff_abs_std_{suffix}\",\n\n        # enmo - anglez\n        f\"enmo_minus_anglez_mean_abs_{suffix}\",\n        f\"enmo_minus_anglez_max_abs_{suffix}\",\n        f\"enmo_minus_anglez_std_abs_{suffix}\",\n    ])","metadata":{"trusted":true,"id":"yKVOIsLYhOJT","execution":{"iopub.status.busy":"2026-05-04T08:30:01.722986Z","iopub.execute_input":"2026-05-04T08:30:01.723541Z","iopub.status.idle":"2026-05-04T08:30:01.733088Z","shell.execute_reply.started":"2026-05-04T08:30:01.723465Z","shell.execute_reply":"2026-05-04T08:30:01.732349Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def to_coord(x, x_max):\n    \"\"\"\n        Adapted from: https://github.com/sakami0000/child-mind-institute-detect-sleep-states-1st-place/blob/main/kami/run/prepare_data.py\n    \"\"\"\n    rad = 2 * np.pi * (x % x_max) / x_max\n    x_sin = np.sin(rad)\n    x_cos = np.cos(rad)\n\n    return x_sin, x_cos\n\ndef extract_window_features(series_id_df, widths, events_exist = True):\n    static_feats = [*to_coord(series_id_df[\"hour\"], 24)]\n    \n    rolling_feats = [\n        feat\n        for mins in widths\n        for feat in [\n            series_id_df[\"enmo\"].rolling(window=12 * mins, center=True, min_periods=1).mean().abs(),\n            series_id_df[\"enmo\"].rolling(window=12 * mins, center=True, min_periods=1).max().abs(),\n            series_id_df[\"enmo\"].rolling(window=12 * mins, center=True, min_periods=1).std().abs(),\n\n            series_id_df[\"enmo\"].diff().abs().rolling(window=12 * mins, center=True, min_periods=1).mean(),\n            series_id_df[\"enmo\"].diff().abs().rolling(window=12 * mins, center=True, min_periods=1).max(),\n            series_id_df[\"enmo\"].diff().abs().rolling(window=12 * mins, center=True, min_periods=1).std(),\n\n            series_id_df[\"anglez\"].rolling(window=12 * mins, center=True, min_periods=1).mean().abs(),\n            series_id_df[\"anglez\"].rolling(window=12 * mins, center=True, min_periods=1).max().abs(),\n            series_id_df[\"anglez\"].rolling(window=12 * mins, center=True, min_periods=1).std().abs(),\n\n            series_id_df[\"anglez\"].diff().abs().rolling(window=12 * mins, center=True, min_periods=1).mean(),\n            series_id_df[\"anglez\"].diff().abs().rolling(window=12 * mins, center=True, min_periods=1).max(),\n            series_id_df[\"anglez\"].diff().abs().rolling(window=12 * mins, center=True, min_periods=1).std(),\n\n            (series_id_df[\"enmo\"] - series_id_df[\"anglez\"]).rolling(window=12 * mins, center=True, min_periods=1).mean().abs(),\n            (series_id_df[\"enmo\"] - series_id_df[\"anglez\"]).rolling(window=12 * mins, center=True, min_periods=1).max().abs(),\n            (series_id_df[\"enmo\"] - series_id_df[\"anglez\"]).rolling(window=12 * mins, center=True, min_periods=1).std().abs(),\n        ]\n    ]\n\n    X = np.column_stack(static_feats + rolling_feats).astype(np.float16)\n\n    if events_exist:\n        y = series_id_df[\"event\"].astype(\"uint8\").to_numpy()\n        return X, y\n    else:\n        return X","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:30:01.734117Z","iopub.execute_input":"2026-05-04T08:30:01.734455Z","iopub.status.idle":"2026-05-04T08:30:01.746383Z","shell.execute_reply.started":"2026-05-04T08:30:01.734432Z","shell.execute_reply":"2026-05-04T08:30:01.745459Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The below loop processes each series_id individually:\n\n* Extracts features using the defined function\n* Aligns labels\n* Saves the resulting feature (`X`) and label (`y`) arrays as `.npy` files\n\nThis approach avoids memory overflow by handling data in chunks.\n\nThe arrays are saved as the dataset [CMI-processed-data](https://www.kaggle.com/datasets/hadighahroudi/cmi-processed-data) and will be reloaded.","metadata":{"id":"pi0Q6PNThOJU"}},{"cell_type":"code","source":"# !rm -r /kaggle/working/*","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:30:01.747396Z","iopub.execute_input":"2026-05-04T08:30:01.747799Z","iopub.status.idle":"2026-05-04T08:30:01.760239Z","shell.execute_reply.started":"2026-05-04T08:30:01.747774Z","shell.execute_reply":"2026-05-04T08:30:01.759550Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not processed_data_exists:\n    os.makedirs(processed_data_dir, exist_ok = True)\n    \n    for series_id in tqdm(series_ids, \"Finding all data item labels and feature extraction...\"):\n        series_id_df = train_set[train_set[\"series_id\"] == series_id].copy()\n        event_series = series_id_df[~series_id_df[\"event\"].isna()].sort_values(\"step\")\n        events = event_series[\"event\"].tolist()\n\n        # Check if there are even number of events (onset + wakeup pairs) and onset events are followed by the corresponding wakeups\n        is_valid = (\n            len(events) % 2 == 0 and\n            all(events[i] == CLASSES[\"onset\"] for i in range(0, len(events), 2)) and\n            all(events[i] == CLASSES[\"wakeup\"] for i in range(1, len(events), 2))\n        )\n\n        if not is_valid:\n            print(f\"Invalid event pattern in series {series_id}: {events}, skipping.\")\n        else:\n            # Assign 0 to all steps initially (default not-sleep)\n            series_id_df[\"event\"] = 0\n\n            for (i, row_onset), (j, row_wakeup) in zip(event_series[::2].iterrows(), event_series[1::2].iterrows()):\n                onset_step = row_onset[\"step\"]\n                wakeup_step = row_wakeup[\"step\"]\n\n                series_id_df.loc[(series_id_df[\"step\"] >= onset_step) & (series_id_df[\"step\"] < wakeup_step), \"event\"] = 1\n\n            X, y = extract_window_features(series_id_df = series_id_df, widths = window_widths)\n\n            data = np.concat([X, y[:, np.newaxis]], axis = 1)\n            np.save(f\"{processed_data_dir}/X_y_{\"-\".join(map(str, window_widths))}m_windows_{series_id}.npy\", data)","metadata":{"trusted":true,"id":"Dm6_y4VLhOJU","execution":{"iopub.status.busy":"2026-05-04T08:30:01.761144Z","iopub.execute_input":"2026-05-04T08:30:01.761678Z","iopub.status.idle":"2026-05-04T08:30:01.772657Z","shell.execute_reply.started":"2026-05-04T08:30:01.761653Z","shell.execute_reply":"2026-05-04T08:30:01.771847Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not os.path.exists(\"X_train.npy\"):\n    data_paths = list(glob(f\"/kaggle/input/cmi-processed-data/{processed_data_dir}/X_y_{\"-\".join(map(str, window_widths))}m_windows_*.npy\"))\n    \n    X_all = []\n    y_all = []\n    for path in data_paths:\n        data = np.load(path).astype(np.float16)\n        X_all.append(data[:, :-1])\n        y_all.append(data[:, -1].squeeze().astype(np.bool))\n    \n    print(len(X_all), len(y_all), X_all[0].shape, y_all[0].shape, X_all[0].dtype, y_all[0].dtype)","metadata":{"trusted":true,"id":"CC4mnJvBhOJU","outputId":"e7124575-8814-49bf-d445-449bef02b994","execution":{"iopub.status.busy":"2026-05-04T08:30:01.773954Z","iopub.execute_input":"2026-05-04T08:30:01.774304Z","iopub.status.idle":"2026-05-04T08:30:01.783033Z","shell.execute_reply.started":"2026-05-04T08:30:01.774280Z","shell.execute_reply":"2026-05-04T08:30:01.782222Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Train-Test-Validation Split\nWe perform splitting at the individual (`series_id`) level.","metadata":{"id":"_uuWSBXwhOJU"}},{"cell_type":"code","source":"if not os.path.exists(\"X_train.npy\"):\n    X_train, X_test, y_train, y_test = train_test_split(X_all, y_all, test_size = 0.2, shuffle = True, random_state = SEED)\n    del X_test, y_test, X_all, y_all\n    gc.collect()\n    \n    X_train, X_valid, y_train, y_valid = train_test_split(X_train, y_train, test_size = 0.1, shuffle = True, random_state = SEED)\n\n    X_train, y_train = np.concat(X_train), np.concat(y_train)\n    X_valid, y_valid = np.concat(X_valid), np.concat(y_valid)\n    \n    print(X_train.shape, y_train.shape, X_valid.shape, y_valid.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:30:01.784000Z","iopub.execute_input":"2026-05-04T08:30:01.784634Z","iopub.status.idle":"2026-05-04T08:30:01.798451Z","shell.execute_reply.started":"2026-05-04T08:30:01.784604Z","shell.execute_reply":"2026-05-04T08:30:01.797451Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# np.save(\"X_train.npy\", X_train)\n# np.save(\"y_train.npy\", y_train)\n# np.save(\"X_valid.npy\", X_valid)\n# np.save(\"y_valid.npy\", y_valid)\n\nif os.path.exists(\"X_train.npy\"):\n    X_train = np.load(\"X_train.npy\")\n    y_train = np.load(\"y_train.npy\")\n    X_valid = np.load(\"X_valid.npy\")\n    y_valid = np.load(\"y_valid.npy\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Feature selection based on the indices determined via SHAP analysis","metadata":{}},{"cell_type":"code","source":"feat_inds_path = \"\" #\"/kaggle/working/shap_top_feats_inds.npy\"\nif os.path.exists(feat_inds_path):\n    top_feats_inds = np.load(feat_inds_path)\n    in_features = len(top_feats_inds)\n    X_train = X_train[:, top_feats_inds]\n    X_valid = X_valid[:, top_feats_inds]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(X_train.shape, X_valid.shape)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. Creating torch dataset and dataloader\nWe create a simple dataset to load data in batches.\n\nTo better leverage the GPU memory and avoid creating a bottleneck, we pre-load all the data to the GPU and then present it to the model in batches.","metadata":{"id":"rO4GL0AehOJV"}},{"cell_type":"code","source":"if norm_type == \"robust\":\n    q1, q2, q3 = np.percentile(X_train, [25, 50, 75], axis = 0)\n    iqr = q3 - q1","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_train = torch.from_numpy(X_train).to(device)\ny_train = torch.from_numpy(y_train).to(device).unsqueeze(-1).type(torch.float16)\nX_valid = torch.from_numpy(X_valid).to(device)\ny_valid = torch.from_numpy(y_valid).to(device).unsqueeze(-1).type(torch.float16)","metadata":{"trusted":true,"id":"oLelizYChOJV"},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Normalization","metadata":{"id":"e6Bp2M-YhOJV"}},{"cell_type":"code","source":"# Z-score normalization\nif norm_type == \"z-score\":\n    train_mean = X_train.mean(dim=0, keepdim=True)\n    train_std = X_train.std(dim=0, keepdim=True)\n    epsilon = 1e-8\n    train_std.add_(epsilon)\n    \n    # Apply in-place normalization to Training Data\n    # .sub_() does in-place subtraction, .div_() does in-place division\n    X_train.sub_(train_mean).div_(train_std)\n    # Apply in-place normalization to Validation Data using the Train stats\n    X_valid.sub_(train_mean).div_(train_std)\n    \n    np.save(\"train_mean.npy\", train_mean.cpu().numpy())\n    np.save(\"train_std.npy\", train_std.cpu().numpy())\n    \n    del train_mean, train_std\n\n# Robust normalization\nif norm_type == \"robust\":\n    q2 = torch.tensor(q2, dtype = X_train.dtype, device = X_train.device)\n    iqr = torch.tensor(iqr, dtype = X_train.dtype, device = X_train.device)\n    X_train.sub_(q2).div_(iqr)\n    X_valid.sub_(q2).div_(iqr)\n\n    del q2, iqr\n    \ntorch.cuda.empty_cache()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(X_train.min(), X_train.max(), X_valid.min(), X_valid.max(), y_train.min(), y_valid.max())","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 8. Building the model ⚙️\nHere we design a conventional MLP with residual connections to allow us increase the number of layers with problems like vanishing gradient.","metadata":{"id":"M8V74cJwhOJV"}},{"cell_type":"code","source":"class ResMLP(nn.Module):\n  def __init__(self, in_features, hidden_features, out_features, dropout = 0):\n    super(ResMLP, self).__init__()\n\n    layers = [\n        nn.Sequential(\n            nn.Linear(in_features = in_features, out_features = hidden_features[0]),\n            nn.BatchNorm1d(hidden_features[0]),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n    ]\n\n    for i in range(len(hidden_features)-1):\n      layers.append(nn.Sequential(\n          nn.Linear(in_features = hidden_features[i], out_features = hidden_features[i+1]),\n          nn.BatchNorm1d(hidden_features[i+1]),\n          nn.ReLU(),\n          nn.Dropout(dropout)\n      ))\n\n    layers.append(nn.Linear(in_features = hidden_features[-1], out_features = out_features))\n\n    self.layers = nn.Sequential(*layers)\n\n  def forward(self, x):\n    x = self.layers[0](x)\n\n    for layer in self.layers[1:-1]:\n      out = layer(x)\n      if x.shape[1] == out.shape[1]:\n        x = x + out\n      else:\n        x = out\n\n    return self.layers[-1](x)\n\n  def predict_shap(self, x): # For using with SHAP\n    x = torch.as_tensor(x, device = next(self.parameters()).device)\n      \n    with torch.inference_mode():\n        logits = self.forward(x)\n        out = torch.exp(logits).detach().cpu().numpy()\n\n    return out\n","metadata":{"trusted":true,"id":"mT0fZGeShOJV","execution":{"iopub.status.busy":"2026-05-04T08:38:51.704655Z","iopub.execute_input":"2026-05-04T08:38:51.705092Z","iopub.status.idle":"2026-05-04T08:38:51.713555Z","shell.execute_reply.started":"2026-05-04T08:38:51.705059Z","shell.execute_reply":"2026-05-04T08:38:51.712705Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9. Training and evaluation\nUtilizing Adam optimizer, we train the model for 50 epochs. We further use automatic mixed precision (AMP) to reduce memory consumption.","metadata":{"id":"nF_8CenFhOJW"}},{"cell_type":"code","source":"def calculate_epoch_metrics(labels, preds, probs, classes):\n    \"\"\"y_true_list/y_pred_list: list of 1D tensors; y_prob_list: list of [B, C] or [B] tensors\"\"\"\n    y_true = np.asarray(labels)\n    y_pred = np.asarray(preds)\n    y_prob = np.asarray(probs)\n\n    acc = accuracy_score(y_true, y_pred)\n    prec, rec, f1, _ = precision_recall_fscore_support(\n        y_true, y_pred, average=\"macro\", zero_division=0\n    )\n    auc = roc_auc_score(y_true, y_prob)\n\n    return {\n        \"accuracy\": acc,\n        \"precision_macro\": prec,\n        \"recall_macro\": rec,\n        \"f1_macro\": f1,\n        \"auc_macro\": auc,\n    }\n\n\ndef train_step(model, X_train, y_train, optimizer, loss_fn, use_amp = False, amp_scaler = None):\n    \"\"\"\n        More info on automatic mixed precision training: https://docs.pytorch.org/tutorials/recipes/recipes/amp_recipe.html\n    \"\"\"\n    model.train()\n\n    if (use_amp == True) and (amp_scaler is not None):\n        # Runs the forward pass under ``autocast`` for mixed precision training.\n        with torch.autocast(device_type = \"cuda\" if X_train.is_cuda else \"cpu\", dtype = torch.float16):\n            logits = model(X_train)\n            loss = loss_fn(logits, y_train)\n        # Scales loss. Calls ``backward()`` on scaled loss to create scaled gradients.\n        amp_scaler.scale(loss).backward()\n        # ``scaler.step()`` first unscales the gradients of the optimizer's assigned parameters.\n        # If these gradients do not contain ``inf``s or ``NaN``s, optimizer.step() is then called,\n        # otherwise, optimizer.step() is skipped.\n        amp_scaler.step(optimizer)\n        # Updates the scale for next iteration.\n        amp_scaler.update()\n    else:\n        logits = model(X_train)\n        loss = loss_fn(logits, y_train)\n        loss.backward()\n        optimizer.step()\n\n    if logits.shape[-1] == 1: # Binary classification\n        probs = torch.sigmoid(logits)\n        preds = torch.round(probs)\n    else:\n        probs = torch.softmax(logits, dim = 1)\n        preds = torch.argmax(probs, dim = 1)\n\n    optimizer.zero_grad(set_to_none=True)\n\n    return preds, probs, loss\n\n\ndef test_step(model, X_test, y_test = None, loss_fn = None, use_amp = False):\n    model.eval()\n    with torch.inference_mode():\n        if use_amp:\n            with torch.autocast(device_type = \"cuda\" if X_test.is_cuda else \"cpu\", dtype = torch.float16):\n                logits = model(X_test)\n        else:\n            logits = model(X_test)\n\n        if logits.shape[-1] == 1: # Binary classification\n            probs = torch.sigmoid(logits)\n            preds = torch.round(probs)\n        else:\n            probs = torch.softmax(logits, dim = 1)\n            preds = torch.argmax(probs, dim = 1)\n\n        if (y_test is not None) and (loss_fn is not None):\n            if logits.shape[1] == 1 and y_test.ndim == 1: # Binary classification\n                y_test = y_test.unsqueeze(1)\n            loss = loss_fn(logits, y_test)\n            return preds, probs, loss\n\n        return preds, probs\n\n\nstart_time = None\ndef start_timer():\n    global start_time\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n        torch.cuda.synchronize()\n    gc.collect()\n\n    start_time = time.time()\n\ndef end_timer_and_print(verbose):\n    if torch.cuda.is_available():\n        torch.cuda.synchronize()\n\n    end_time = time.time()\n    if verbose:\n        print(\"Total execution time = {:.3f} sec\".format(end_time - start_time))\n        print(\"Max memory used by tensors = {} bytes\".format(torch.cuda.max_memory_allocated()))","metadata":{"trusted":true,"id":"m2fXKC8-hOJW","execution":{"iopub.status.busy":"2026-05-04T08:38:54.971438Z","iopub.execute_input":"2026-05-04T08:38:54.972277Z","iopub.status.idle":"2026-05-04T08:38:54.984346Z","shell.execute_reply.started":"2026-05-04T08:38:54.972246Z","shell.execute_reply":"2026-05-04T08:38:54.983435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_and_validate(model, optim, loss_fn, batch_size, epochs, use_amp,\n                       X_train, y_train, X_valid, y_valid, \n                       augment = False, balance = True, verbose = 1, save_checkpoints = True, debug = False):\n    start_timer()\n    \n    history = []\n    lowest_loss = 1e6\n    \n    scaler = torch.amp.GradScaler(device)\n    \n    # Get and batch the indices of valid samples\n    inds_valid = torch.arange(len(X_valid))\n    inds_valid_batches = torch.split(inds_valid, batch_size)\n    \n    # Find the indices of positive and negative samples\n    inds_train_pos = torch.where(y_train.cpu() == 1)[0]\n    inds_train_neg = torch.where(y_train.cpu() == 0)[0]\n    n_pos = len(inds_train_pos)\n    \n    train_preds, train_probs, train_labels = [], [], []\n    test_preds, test_probs, test_labels = [], [], []\n    \n    for e in range(epochs):\n        train_losses, test_losses = [], []\n    \n        if balance:\n            # Class balancing via random undersampling\n            inds_train_neg_sampled = inds_train_neg[torch.randperm(len(inds_train_neg))[:n_pos]] # Permute and sample the negative sample indices\n            inds_train = torch.cat([inds_train_pos, inds_train_neg_sampled]) # Merge the balanced positive and negative indices\n            inds_train = inds_train[torch.randperm(len(inds_train))] # Permute the merged train indices\n        else:\n            inds_train = torch.randperm(len(X_train))\n        inds_train_batches = torch.split(inds_train, batch_size)\n    \n        if debug:\n            i = 0\n    \n        # ---------- Train loop ----------\n        for batch_inds in tqdm(inds_train_batches, desc=f\"Epoch {e} training...\"):\n            X = X_train[batch_inds]\n            y = y_train[batch_inds]\n\n            if augment:\n                X = X + torch.randn_like(X) * 0.01 # Adding small amount of Gaussian noise for data augmentation\n                \n            pred, prob, loss = train_step(model=model, X_train= X, y_train= y,\n                                          optimizer=optim, loss_fn=loss_fn,\n                                          use_amp = use_amp, amp_scaler = scaler)\n    \n            # Collect (move to CPU later in the helper)\n            if e == epochs - 1:\n                train_labels.append(y.detach().cpu())\n                train_preds.append(pred.detach().cpu())\n                train_probs.append(prob.detach().cpu())\n            train_losses.append(loss.detach().item())\n    \n            if debug:\n                i += 1\n                if i == 10:\n                    break\n    \n        # ---------- Test loop ----------\n        for batch_inds in tqdm(inds_valid_batches, desc=f\"Epoch {e} validation...\"):\n            X = X_valid[batch_inds]\n            y = y_valid[batch_inds]\n            pred, prob, loss = test_step(model = model, X_test = X, y_test = y, loss_fn = loss_fn, use_amp = use_amp)\n    \n            if e == epochs - 1:\n                test_labels.append(y.detach().cpu())\n                test_preds.append(pred.detach().cpu())\n                test_probs.append(prob.detach().cpu())\n            test_losses.append(loss.detach().item())\n    \n            if debug:\n                i += 1\n                if i == 20:\n                    break\n    \n        # ---------- Logging ----------\n        if verbose:\n            print(f\"Epoch [{e}/{epochs}] | train loss={np.mean(train_losses):.4f} | test loss={np.mean(test_losses):.4f}\")\n    \n        history.append({\n            \"epoch\": e,\n            \"train_loss\": float(np.mean(train_losses)),\n            \"test_loss\": float(np.mean(test_losses))\n        })\n\n    \n        if save_checkpoints:\n            torch.save(model.state_dict(), \"model_checkpoints/ResMLP_last.pth\")\n            if history[-1][\"test_loss\"] < lowest_loss:\n                lowest_loss = history[-1][\"test_loss\"]\n                torch.save(model.state_dict(), \"model_checkpoints/ResMLP_lowest_loss.pth\")\n            \n    end_timer_and_print(verbose)\n    \n    train_metrics = calculate_epoch_metrics(torch.concat(train_labels), torch.concat(train_preds), torch.concat(train_probs), list(range(len(CLASSES))))\n    test_metrics = calculate_epoch_metrics(torch.concat(test_labels), torch.concat(test_preds), torch.concat(test_probs), list(range(len(CLASSES))))\n\n    return history, train_metrics, test_metrics","metadata":{"trusted":true,"id":"Zt3wEfyIhOJW","outputId":"43edec25-02eb-49c5-c0ed-8482eb76a721"},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 9.1. Hyperparameter tuning\nBalance the data and sample 10_000_000 records for grid search:","metadata":{}},{"cell_type":"code","source":"# inds_train_pos = torch.where(y_train == 1)[0]\n# inds_train_neg = torch.where(y_train == 0)[0]\n# n_pos = len(inds_train_pos)\n\n# inds_train_neg_sampled = inds_train_neg[torch.randperm(len(inds_train_neg))[:n_pos]] # Permute and sample the negative sample indices\n# inds_train = torch.cat([inds_train_pos[:5_000_000], inds_train_neg_sampled[:5_000_000]]) # Merge the balanced positive and negative indices\n# inds_train = inds_train[torch.randperm(len(inds_train))] # Permute the merged train indices\n# X_train_sample = X_train[inds_train]\n# y_train_sample = y_train[inds_train]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# def objective(trial):\n#     # Hyperparameters to tune\n#     n_layers = trial.suggest_int(\"n_layers\", 2, 7)\n#     hidden_features = [\n#         trial.suggest_categorical(f\"n_units_l{i}\", [16, 32, 64, 128, 256, 512, 1024])\n#         for i in range(n_layers)\n#     ]\n\n#     learning_rate = trial.suggest_categorical(\"lr\", [0.001, 0.0001, 0.00001])\n#     batch_size = trial.suggest_categorical(\"batch_size\", [1000, 10000, 100000])\n#     epochs = 5\n#     use_amp = True\n\n#     model = ResMLP(\n#         in_features = in_features,\n#         hidden_features = hidden_features,\n#         out_features = out_features\n#     ).to(device)\n    \n#     optim = torch.optim.Adam(model.parameters(), lr = learning_rate, weight_decay = weight_decay)\n#     loss_fn = nn.BCEWithLogitsLoss()\n\n#     history, train_metrics, test_metrics = train_and_validate(model = model, optim = optim, loss_fn = loss_fn, batch_size = batch_size, epochs = epochs, use_amp = use_amp,\n#                                                              X_train = X_train_sample, y_train = y_train_sample, X_valid = X_valid, y_valid = y_valid, \n#                                                              balance_and_permute = False, verbose = 0, save_checkpoints = False)\n\n#     return test_metrics[\"f1_macro\"]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # Run Study\n# study = optuna.create_study(direction=\"maximize\")\n# study.optimize(objective, n_trials=20)\n\n# # Best Results\n# print(\"Best trial:\")\n# trial = study.best_trial\n\n# print(\"  f1_macro:\", trial.value)\n# print(\"  Params:\")\n# for key, value in trial.params.items():\n#     print(f\"    {key}: {value}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 9.2. Model fitting and validation","metadata":{}},{"cell_type":"code","source":"# hidden_features = [trial.params[f\"n_units_l{i}\"] for i in range(trial.params[\"n_layers\"])]\n# learning_rate = trial.params[\"lr\"]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = ResMLP(in_features=in_features, hidden_features=hidden_features, out_features=out_features, dropout = 0.3).to(device)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"optim = torch.optim.Adam(model.parameters(), lr = learning_rate, weight_decay = weight_decay)\nloss_fn = nn.BCEWithLogitsLoss()\n\nos.makedirs(\"model_checkpoints\", exist_ok = True)","metadata":{"trusted":true,"id":"B8NkShEYhOJW"},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# del inds_train_pos, inds_train_neg, inds_train_neg_sampled, inds_train, X_train_sample, y_train_sample\n# gc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"history, train_metrics, test_metrics = train_and_validate(model = model, optim = optim, loss_fn = loss_fn, batch_size = batch_size, epochs = epochs, use_amp = True,\n                                                         X_train = X_train, y_train = y_train, X_valid = X_valid, y_valid = y_valid, augment = True)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\n    f\"train acc={train_metrics['accuracy']:.4f}, prec={train_metrics['precision_macro']:.4f}, \"\n    f\"rec={train_metrics['recall_macro']:.4f}, f1={train_metrics['f1_macro']:.4f}, auc={train_metrics['auc_macro']:.4f} \"\n    f\"| test acc={test_metrics['accuracy']:.4f}, prec={test_metrics['precision_macro']:.4f}, \"\n    f\"rec={test_metrics['recall_macro']:.4f}, f1={test_metrics['f1_macro']:.4f}, auc={test_metrics['auc_macro']:.4f}\"\n)","metadata":{"trusted":true,"id":"MzM0iy4LhOJW","outputId":"818e79bd-41a4-4bab-8f22-2120202976de"},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\n# Extract data\nlst_epochs = [h[\"epoch\"] for h in history]\ntrain_loss = [h[\"train_loss\"] for h in history]\ntest_loss = [h[\"test_loss\"] for h in history]\n\n# Global style (clean academic look)\nplt.style.use(\"seaborn-v0_8-whitegrid\")\n\nfig, ax = plt.subplots(figsize=(9, 5.5), dpi=120)\n\nax.plot(lst_epochs, train_loss, label=\"Train Loss\", linewidth=2)\nax.plot(lst_epochs, test_loss, label=\"Test Loss\", linewidth=2)\n\nax.set_title(\"Training and Validation Loss Over Epochs\", fontsize=14, pad=12)\nax.set_xlabel(\"Epoch\", fontsize=12)\nax.set_ylabel(\"Loss\", fontsize=12)\n\nax.legend(frameon=True)\nax.spines[\"top\"].set_visible(False)\nax.spines[\"right\"].set_visible(False)\n\nax.tick_params(axis=\"both\", labelsize=10)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 10. Features analysis","metadata":{}},{"cell_type":"code","source":"# np.save(\"X_train_norm.npy\", X_train)\n# np.save(\"y_train.npy\", y_train)\n# np.save(\"X_valid_norm.npy\", X_valid)\n# np.save(\"y_valid.npy\", y_valid)\n\n# X_train = np.load(\"X_train_norm.npy\")\n# y_train = np.load(\"y_train.npy\")\n# X_valid = np.load(\"X_valid_norm.npy\")\n# y_valid = np.load(\"y_valid.npy\")\n\n# X_train = torch.from_numpy(X_train).to(device)\n# y_train = torch.from_numpy(y_train).to(device).unsqueeze(-1).type(torch.float16)\n# X_valid = torch.from_numpy(X_valid).to(device)\n# y_valid = torch.from_numpy(y_valid).to(device).unsqueeze(-1).type(torch.float16)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# model.load_state_dict(torch.load(\"/kaggle/working/model_checkpoints/ResMLP_lowest_loss.pth\"))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"background = X_train[torch.randperm(X_train.shape[0])[:100]].cpu().type(torch.float32).numpy()  # small subset\nX = X_valid[torch.randperm(X_valid.shape[0])[:5000]]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\nexplainer = shap.KernelExplainer(model.predict_shap, background)\nshap_values = explainer.shap_values(X.cpu().type(torch.float32).numpy())","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"shap.summary_plot(shap_values = shap_values.squeeze(), features = X, feature_names = feature_names)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"> When we are explaining a prediction 𝑓⁡(𝑥), the SHAP value for a specific feature 𝑖 is just the difference between the expected model output and the partial dependence plot at the feature’s value 𝑥𝑖:","metadata":{}},{"cell_type":"code","source":"# make a standard partial dependence plot\nshap.partial_dependence_plot(\n    0,\n    model.predict_shap,\n    background,\n    feature_names = feature_names,\n    ice=False,\n    model_expected_value=True,\n    feature_expected_value=True\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"feats_importance = np.mean(np.abs(shap_values), axis = 0).squeeze()\n\nk = 32\ntop_k_ids = np.argsort(feats_importance)[::-1][:k]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for i in top_k_ids:\n    print(i, feature_names[i], feats_importance[i])","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# np.save(\"shap_top_feats_inds.npy\", top_k_ids)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 11. Retraining on All Data and Internal Testing","metadata":{}},{"cell_type":"code","source":"data_paths = list(glob(f\"/kaggle/input/cmi-processed-data/{processed_data_dir}/X_y_{\"-\".join(map(str, window_widths))}m_windows_*.npy\"))\n\nX_all = []\ny_all = []\nfor path in data_paths:\n    data = np.load(path).astype(np.float16)\n    X_all.append(data[:, :-1])\n    y_all.append(data[:, -1].squeeze().astype(np.bool))\n\nprint(len(X_all), len(y_all), X_all[0].shape, y_all[0].shape, X_all[0].dtype, y_all[0].dtype)\n\nX_train, X_test, y_train, y_test = train_test_split(X_all, y_all, test_size = 0.2, shuffle = True, random_state = SEED)\n\ndel X_train, y_train, X_all, y_all, data\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:32:02.752325Z","iopub.execute_input":"2026-05-04T08:32:02.753077Z","iopub.status.idle":"2026-05-04T08:33:44.678563Z","shell.execute_reply.started":"2026-05-04T08:32:02.753044Z","shell.execute_reply":"2026-05-04T08:33:44.677550Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_test, y_test = np.concat(X_test), np.concat(y_test)\n\nX_test = torch.from_numpy(X_test).to(device)\ny_test = torch.from_numpy(y_test).to(device).unsqueeze(-1).type(torch.float16)\n\nprint(X_test.shape, y_test.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:35:26.715248Z","iopub.execute_input":"2026-05-04T08:35:26.715723Z","iopub.status.idle":"2026-05-04T08:35:27.890944Z","shell.execute_reply.started":"2026-05-04T08:35:26.715691Z","shell.execute_reply":"2026-05-04T08:35:27.890161Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_mean = torch.from_numpy(np.load(\"train_mean.npy\")).to(device)\ntrain_std = torch.from_numpy(np.load(\"train_std.npy\")).to(device)\nX_test.sub_(train_mean).div_(train_std)\n\nprint(X_test.min(), X_test.max(), y_test.min(), y_test.max())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:36:52.056538Z","iopub.execute_input":"2026-05-04T08:36:52.057427Z","iopub.status.idle":"2026-05-04T08:36:52.328110Z","shell.execute_reply.started":"2026-05-04T08:36:52.057384Z","shell.execute_reply":"2026-05-04T08:36:52.327235Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = ResMLP(in_features=in_features, hidden_features=hidden_features, out_features=out_features, dropout = 0.3).to(device)\nmodel.load_state_dict(torch.load(\"/kaggle/working/model_checkpoints/ResMLP_last.pth\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:44:29.589966Z","iopub.execute_input":"2026-05-04T08:44:29.590953Z","iopub.status.idle":"2026-05-04T08:44:29.610324Z","shell.execute_reply.started":"2026-05-04T08:44:29.590917Z","shell.execute_reply":"2026-05-04T08:44:29.609661Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pred, prob = test_step(model = model, X_test = X_test, y_test = y_test, use_amp = True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:44:30.307742Z","iopub.execute_input":"2026-05-04T08:44:30.308691Z","iopub.status.idle":"2026-05-04T08:44:30.314907Z","shell.execute_reply.started":"2026-05-04T08:44:30.308656Z","shell.execute_reply":"2026-05-04T08:44:30.314244Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_metrics = calculate_epoch_metrics(y_test.detach().cpu(), pred.detach().cpu(), prob.detach().cpu(), list(range(len(CLASSES))))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:44:31.045312Z","iopub.execute_input":"2026-05-04T08:44:31.046043Z","iopub.status.idle":"2026-05-04T08:44:58.252163Z","shell.execute_reply.started":"2026-05-04T08:44:31.046009Z","shell.execute_reply":"2026-05-04T08:44:58.251497Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\n    f\"train acc={test_metrics['accuracy']:.4f}, prec={test_metrics['precision_macro']:.4f}, \"\n    f\"rec={test_metrics['recall_macro']:.4f}, f1={test_metrics['f1_macro']:.4f}, auc={test_metrics['auc_macro']:.4f} \"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-04T08:44:58.253286Z","iopub.execute_input":"2026-05-04T08:44:58.254086Z","iopub.status.idle":"2026-05-04T08:44:58.258295Z","shell.execute_reply.started":"2026-05-04T08:44:58.254056Z","shell.execute_reply":"2026-05-04T08:44:58.257639Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 12. External testing on the competition data","metadata":{}},{"cell_type":"code","source":"data_paths = list(glob(f\"/kaggle/input/datasets/hadighahroudi/cmi-processed-data/processed_CMI_data/X_y_{\"-\".join(map(str, window_widths))}m_windows_*.npy\"))\n\nX_all = []\ny_all = []\nfor path in data_paths:\n    data = np.load(path).astype(np.float16)\n    X_all.append(data[:, :-1])\n    y_all.append(data[:, -1].squeeze().astype(np.bool))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_all = np.concat(X_all)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"feats_min = X_all.min(axis = 0)\nfeats_max = X_all.max(axis = 0)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del X_all\ngc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_series = pd.read_parquet(\"/kaggle/input/competitions/child-mind-institute-detect-sleep-states/test_series.parquet\").sort_values([\"series_id\", \"step\"])\ntest_series_ids = test_series[\"series_id\"].unique()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_test = []\n\nfor series_id in tqdm(test_series_ids, \"Feature extraction...\"):\n    series_id_df = test_series[test_series[\"series_id\"] == series_id].copy().sort_values(\"step\")\n    X = extract_window_features(series_id_df = series_id_df, widths = window_widths, events_exist = False)\n    X_test.append(X)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_test = torch.from_numpy(np.concat(X_test)).to(device)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_test = (X_test - feats_min) / (feats_max - feats_min)\nprint(X_test.shape, X_test.min(), X_test.max(), X_test.dtype)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = ResMLP(in_features=in_features, hidden_features=hidden_features, out_features=out_features).to(device)\nmodel.load_state_dict(torch.load(\"/kaggle/input/models/hadighahroudi/cmi-sleep-detector/pytorch/resmlp/1/ResMLP_last.pth\", map_location = device))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\nwith torch.inference_mode():\n    with torch.autocast(device_type = \"cuda\" if X_test.is_cuda else \"cpu\", dtype = torch.float16):\n        logits = model(X_test)\n        probs = torch.sigmoid(logits).squeeze()\n        preds = torch.round(probs).squeeze()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pred_onsets = (preds.diff() > 0)\npred_wakeups = (preds.diff() < 0)\npred_onset_inds = torch.where(pred_onsets)[0].tolist()\npred_wakeup_inds = torch.where(pred_wakeups)[0].tolist()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"prob_onsets = probs[pred_onset_inds].tolist()\nprob_wakeups = probs[pred_wakeup_inds].tolist()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_submission = pd.read_csv(\"/kaggle/input/competitions/child-mind-institute-detect-sleep-states/sample_submission.csv\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_result = {\"series_id\": [], \"step\": [], \"event\": [],\"score\": []}\n\nfor ind, score in zip(pred_onset_inds, prob_onsets):\n    series_id = test_series.iloc[ind][\"series_id\"]\n    \n    # Find the beginning of the series_id records and offset the global ind\n    series_id_start_ind = test_series[test_series[\"series_id\"] == series_id].index[0]\n    step = ind - series_id_start_ind\n\n    event = \"onset\"\n\n    final_result[\"series_id\"].append(series_id)\n    final_result[\"step\"].append(step)\n    final_result[\"event\"].append(event)\n    final_result[\"score\"].append(score)\n\n\nfor ind, score in zip(pred_wakeup_inds, prob_wakeups):\n    series_id = test_series.iloc[ind][\"series_id\"]\n    \n    # Find the beginning of the series_id records and offset the global ind\n    series_id_start_ind = test_series[test_series[\"series_id\"] == series_id].index[0]\n    step = ind - series_id_start_ind\n\n    event = \"wakeup\"\n\n    final_result[\"series_id\"].append(series_id)\n    final_result[\"step\"].append(step)\n    final_result[\"event\"].append(event)\n    final_result[\"score\"].append(score)\n    ","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_submission = pd.DataFrame(final_result)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_submission.reset_index(names=\"row_id\").to_csv('submission.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_submission","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_series.loc[pred_onset_inds, \"event\"]= \"onset\"\ntest_series.loc[pred_wakeup_inds, \"event\"] = \"wakeup\"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for series_id in test_series_ids:\n    series_id_df = test_series[test_series[\"series_id\"] == series_id]\n    plt.figure(figsize= (24, 6))\n\n    plt.subplot(211)\n    plt.plot(series_id_df[\"step\"], series_id_df[\"enmo\"])\n    ymin, ymax = plt.ylim()\n    plt.vlines(series_id_df[series_id_df[\"event\"] == \"onset\"][\"step\"], ymin = ymin, ymax = ymax, colors= \"red\", linestyles=\"--\")\n    plt.vlines(series_id_df[series_id_df[\"event\"] == \"wakeup\"][\"step\"], ymin = ymin, ymax = ymax, colors= \"orange\", linestyles=\"--\")\n    plt.title(\"ENMO\")\n\n    plt.subplot(212)\n    plt.plot(series_id_df[\"step\"], series_id_df[\"anglez\"])\n    ymin, ymax = plt.ylim()\n    plt.vlines(series_id_df[series_id_df[\"event\"] == \"onset\"][\"step\"], ymin = ymin, ymax = ymax, colors= \"red\", linestyles=\"--\")\n    plt.vlines(series_id_df[series_id_df[\"event\"] == \"wakeup\"][\"step\"], ymin = ymin, ymax = ymax, colors= \"orange\", linestyles=\"--\")\n    plt.title(\"Z-angle\")\n\n    plt.tight_layout()\n    plt.show()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}