# Training a Neural Network to Estimate Sensitivity [![Open in GitHub](https://img.shields.io/badge/Open-GitHub-black?logo=github)][REPRODUCIBILITY_LINK] GWKokab has easy to build and train APIs to train a Multilayer Perceptron (MLP) to estimate the probability of detection or sensitive spacetime volume. This notebook demonstrates how to use these APIs on toy data. Let's first create some toy data where the target is simply the sum of the input features. ```python import h5py import numpy as np from gwkokab.utils.train import train_regressor data = np.random.normal(size=(10_000, 5)) target = np.sum(data, axis=1, keepdims=True) data_path = "data.hdf5" with h5py.File(data_path, "w") as f: f.create_dataset("x0", data=data[:, 0]) f.create_dataset("x1", data=data[:, 1]) f.create_dataset("x2", data=data[:, 2]) f.create_dataset("x3", data=data[:, 3]) f.create_dataset("x4", data=data[:, 4]) f.create_dataset("y", data=target) ``` Now we can use the [`train_regressor`](https://gwkokab.readthedocs.io/en/latest/autoapi/gwkokab/utils/train/index.html#gwkokab.utils.train.train_regressor) function to train an MLP on this data. We will specify the input and output keys, the architecture of the MLP (width and depth), batch size, data path, checkpoint path, number of epochs, validation split, and learning rate. ```python input_keys = ["x" + str(i) for i in range(5)] output_keys = ["y"] checkpoint_path = "model_checkpoint.hdf5" train_regressor( input_keys=input_keys, output_keys=output_keys, width_size=32, depth=2, batch_size=64, data_path=data_path, checkpoint_path=checkpoint_path, epochs=500, validation_split=0.2, # 20% of data used for validation learning_rate=1e-3, ) ``` ```{toggle} ```md Mon Nov 3 02:33:45 2025 +-----------------------------------------------------------------------------------------+ | NVIDIA-SMI 575.57.08 Driver Version: 575.57.08 CUDA Version: 12.9 | |-----------------------------------------+------------------------+----------------------+ | GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. | | | | MIG M. | |=========================================+========================+======================| | 0 NVIDIA GeForce RTX 4050 ... On | 00000000:01:00.0 Off | N/A | | N/A 37C P0 12W / 30W | 15MiB / 6141MiB | 0% Default | | | | N/A | +-----------------------------------------+------------------------+----------------------+ +-----------------------------------------------------------------------------------------+ | Processes: | | GPU GI CI PID Type Process name GPU Memory | | ID ID Usage | |=========================================================================================| | 0 N/A N/A 1745 G /usr/lib/xorg/Xorg 4MiB | +-----------------------------------------------------------------------------------------+ ``` ```{toggle} ```md 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:251 - Input Keys: x0, x1, x2, x3, x4 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:252 - Output Keys: y 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:253 - Width Size: 32 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:254 - Depth: 2 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:255 - Data Path: data.hdf5 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:256 - Checkpoint Path: model_checkpoint.hdf5 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:257 - Train Size: 8000 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:258 - Test Size: 2000 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:259 - Validation Split: 0.2 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:260 - Batch Size: 64 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:261 - Epochs: 500 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:262 - Learning Rate (peak): 0.001 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:263 - Loss Type: mse 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:264 - Scheduler: cosine+warmup 2025-11-03 01:41:40.050 | INFO | gwkokab.utils.train:train_regressor:265 - Weight Decay: 0.0001, Grad Clip: 1.0, Seed: 42 Epoch 1/500, Loss: 4.78830E+00, Val Loss: 4.41401E+00: 0%| | 0/500 [00:06 All the code and files used in this tutorial can be found in [hello-gwkokab/training_mlp][REPRODUCIBILITY_LINK]. [REPRODUCIBILITY_LINK]: https://github.com/kokabsc/hello-gwkokab/tree/main/training_mlp