From 06af4b9cc4e5f01180a4f6f77ade8efab569e597 Mon Sep 17 00:00:00 2001 From: Asger Gitz-Johansen Date: Sun, 2 Aug 2026 15:01:22 +0200 Subject: [PATCH] feat: much better architecture I understand at least some of this. --- src/sleep_detection/main.py | 234 +++++++++++------------------------- 1 file changed, 69 insertions(+), 165 deletions(-) diff --git a/src/sleep_detection/main.py b/src/sleep_detection/main.py index 30b8c35..b08204d 100644 --- a/src/sleep_detection/main.py +++ b/src/sleep_detection/main.py @@ -3,6 +3,7 @@ import sys from argparse import ArgumentParser from argparse import Namespace +from collections import Counter from collections.abc import Sequence from pathlib import Path @@ -17,125 +18,39 @@ from torch.optim import SGD from sleep_detection.sleep_epoch_classifer import SleepEpochClassifier -SLEEP_STAGE_LABELS = [0, 1, 2] - - -def normalize_sleep_stage(stage) -> int: - """Map raw sleep stages to the 3-class target space. - - Raw labels: - 0 -> 0 - 1, 2, 3, 4 -> 1 - 5, 6 -> 2 - Other values are treated as unknown. - """ - if pandas.isna(stage): - return int(numpy.nan) - - value = int(stage) - if value == 0: - return 0 - if 1 <= value <= 4: - return 1 - if 5 <= value <= 6: - return 2 - - return int(numpy.nan) - @logger.catch def main() -> None: """The Main entrypoint function.""" cli = parse_args() - logger.info("welcome") - load_session(cli.a, cli.H, cli.l) - raise ValueError(f"Unknown command: {cli.command}") + session = load_session(cli.a, cli.H, cli.l) + + inputs = torch.tensor(session[["timestamp", "heartrate", "motion-avg"]].to_numpy(), dtype=torch.float32) + mean = inputs.mean(dim=0, keepdim=True) # shape (1, n_features) + std = inputs.std(dim=0, keepdim=True) + 1e-8 + normalized_inputs = (inputs - mean) / std + logger.info(f"inputs={normalized_inputs}") + + targets = torch.tensor(session["stage"].to_numpy(), dtype=torch.long) + logger.info(f"targets={targets}") + + net = SleepEpochClassifier(n_features=3, hidden_size=16, n_classes=3) + train(normalized_inputs, targets, net) + torch.save(net.state_dict(), cli.output) def parse_args() -> Namespace: + """Parse commandline arguments.""" parser = ArgumentParser(description="Training program.") parser.add_argument("-v", "--verbose", help="Increase verbosity.", action="store_true") parser.add_argument("-V", "--version", action="store_true", help="Show version") - parser.add_argument("-a", type=Path, required=True, help="Acceleration data file(s).") parser.add_argument("-H", type=Path, required=True, help="Heartrate data file(s).") parser.add_argument("-l", type=Path, required=True, help="Labels data file(s).") - + parser.add_argument("-o", "--output", type=Path, required=True, help="Output .pt file.") return parser.parse_args() -def featureize(data: pandas.DataFrame) -> tuple[pandas.DataFrame, Tensor, Tensor, list]: - """Convert loaded data into input and target tensors. - - Args: - data: Loaded data, should contain the keys: [timestamp, heartrate, motion, x, y, z, stage] - - """ - logger.info(f"Featurizing the data...") - epochs = build_epochs(data) - - epochs["hr_std"] = epochs["hr_std"].fillna(0.0) - epochs = epochs.dropna(subset=["hr_mean", "motion_mean", "motion_max", "stage"]) - - epochs["seconds_since_start"] = epochs.index.total_seconds() - - CONTEXT = 2 # epochs each side -> 5 epochs total = 25 min window - FEATURE_COLS = ["hr_mean", "hr_std", "motion_mean", "motion_max", "seconds_since_start"] - - features = epochs[FEATURE_COLS].to_numpy(dtype=numpy.float32) - labels = epochs["stage"].astype(numpy.int64).to_numpy() - - X, y = [], [] - for i in range(CONTEXT, len(epochs) - CONTEXT): - window = features[i - CONTEXT : i + CONTEXT + 1] # shape: (2*CONTEXT+1, n_feats) - X.append(window.flatten()) # -> 1D vector - y.append(labels[i]) - - if not X: - raise ValueError("Not enough complete epochs to build a training set.") - - X = torch.tensor(numpy.stack(X), dtype=torch.float32) # shape: (n_epochs, (2*CONTEXT+1)*n_feats) - y = torch.tensor(y, dtype=torch.long) # shape: (n_epochs,) - return epochs, X, y, SLEEP_STAGE_LABELS - - -def featureize_sessions(dataframes: Sequence[pandas.DataFrame]) -> tuple[pandas.DataFrame, Tensor, Tensor, list]: - """Convert multiple loaded sessions into input and target tensors.""" - epochs_per_session = [] - for data in dataframes: - epochs = build_epochs(data) - - epochs["hr_std"] = epochs["hr_std"].fillna(0.0) - epochs = epochs.dropna(subset=["hr_mean", "motion_mean", "motion_max", "stage"]) - epochs["seconds_since_start"] = epochs.index.total_seconds() - epochs_per_session.append(epochs) - - if not epochs_per_session: - raise ValueError("No sessions provided.") - - combined_epochs = pandas.concat(epochs_per_session) - - CONTEXT = 2 # epochs each side -> 5 epochs total = 25 min window - FEATURE_COLS = ["hr_mean", "hr_std", "motion_mean", "motion_max", "seconds_since_start"] - - X, y = [], [] - for epochs in epochs_per_session: - features = epochs[FEATURE_COLS].to_numpy(dtype=numpy.float32) - labels = epochs["stage"].astype(numpy.int64).to_numpy() - - for i in range(CONTEXT, len(epochs) - CONTEXT): - window = features[i - CONTEXT : i + CONTEXT + 1] # shape: (2*CONTEXT+1, n_feats) - X.append(window.flatten()) # -> 1D vector - y.append(labels[i]) - - if not X: - raise ValueError("Not enough complete epochs to build a training set.") - - X = torch.tensor(numpy.stack(X), dtype=torch.float32) # shape: (n_epochs, (2*CONTEXT+1)*n_feats) - y = torch.tensor(y, dtype=torch.long) # shape: (n_epochs,) - return combined_epochs, X, y, SLEEP_STAGE_LABELS - - def load_session(acceleration_file: Path, heartrate_file: Path, label_file: Path) -> pandas.DataFrame: """Load a session triplet.""" adf = load_acceleration_data(acceleration_file) @@ -152,48 +67,10 @@ def load_session(acceleration_file: Path, heartrate_file: Path, label_file: Path # Label the heartrate dataset. ldf.index = ldf.index.astype(float) hdf = pandas.merge_asof(hdf, ldf, on="timestamp", direction="backward") - with pandas.option_context('display.max_rows', None, 'display.max_columns', None): - print(f"{hdf}") logger.info(f"\n{hdf}") return hdf -def build_epochs(data: pandas.DataFrame) -> pandas.DataFrame: - """Aggregate raw signals into 5-minute epochs.""" - EPOCH_LEN = pandas.Timedelta(minutes=5) - - return data.resample(EPOCH_LEN).agg( - hr_mean=("hr", "mean"), - hr_std=("hr", "std"), - motion_mean=("motion", "mean"), - motion_max=("motion", "max"), - stage=("stage", lambda s: s.dropna().mode().iloc[0] if len(s.dropna()) else numpy.nan), - ) - - -def execute(epochs: pandas.DataFrame, input: Tensor, target: Tensor, classifier: Module, label_mapping: list) -> None: - """Run the model and emit per-epoch predictions.""" - classifier.eval() - with torch.no_grad(): - output = classifier(input) - predicted = output.argmax(dim=1) - - decoded_predicted = [label_mapping[index] for index in predicted.tolist()] - decoded_target = [label_mapping[index] for index in target.tolist()] - - report = pandas.DataFrame( - { - "timestamp": epochs.index[2 : len(epochs) - 2].astype("timedelta64[s]") / numpy.timedelta64(1, "s"), - "predicted": decoded_predicted, - "actual": decoded_target, - } - ) - report.to_csv(sys.stdout, index=False) - - accuracy = (predicted == target).float().mean().item() - logger.info(f"accuracy={accuracy:.3f}") - - def load_acceleration_data(data_file: Path) -> pandas.DataFrame: """Load acceleration data file.""" logger.info(f"Loading acceleration data file {data_file}") @@ -219,49 +96,76 @@ def load_labels_data(data_file: Path) -> pandas.DataFrame: """Load label data file.""" logger.info(f"Loading label date file {data_file}") df = pandas.read_csv(data_file, sep=r"\s+", names=["timestamp", "stage"], header=None) - # df["stage"] = pandas.to_numeric(df["stage"], errors="coerce").map(normalize_sleep_stage) + df["stage"] = pandas.to_numeric(df["stage"], errors="coerce").map(normalize_sleep_stage) df = df.drop(df[df["timestamp"] < 0].index).drop_duplicates() df = df.set_index("timestamp") logger.info(f"({len(df)})\n{df}") return df -def train(input: Tensor, target: Tensor, classifier: Module): - classifier.train() - optimizer = SGD(classifier.parameters(), lr=0.01) - criterion = CrossEntropyLoss() +def normalize_sleep_stage(stage) -> int: + """Map raw sleep stages to the 3-class target space.""" + if pandas.isna(stage): + return int(numpy.nan) + value = int(stage) + if value == 0: + return 0 + if 1 <= value <= 4: + return 1 + if 5 <= value <= 6: + return 2 + return int(numpy.nan) + +def train(input: Tensor, target: Tensor, net: Module) -> None: + net.train() + optimizer = SGD(net.parameters(), lr=0.001, momentum=0.9) + + # 20/80 split the input so we can get an accurate `accuracy` log. n_samples = input.shape[0] n_train = max(1, int(n_samples * 0.8)) permutation = torch.randperm(n_samples) - train_idx = permutation[:n_train] - val_idx = permutation[n_train:] + train_idx, val_idx = permutation[:n_train], permutation[n_train:] + train_input, train_target = input[train_idx], target[train_idx] + val_input, val_target = input[val_idx], target[val_idx] - train_input = input[train_idx] - train_target = target[train_idx] - val_input = input[val_idx] if len(val_idx) else None - val_target = target[val_idx] if len(val_idx) else None + # Tweak the loss function. + class_counts = torch.bincount(train_target) + class_weights = 1.0 / torch.sqrt(class_counts.float()) + class_weights = class_weights / class_weights.sum() # normalize, not strictly required but keeps scale sane + criterion = CrossEntropyLoss(weight=class_weights) + # Train for a bunch of epochs. for epoch in range(100): optimizer.zero_grad() - output = classifier(train_input) + output = net(train_input) loss = criterion(output, train_target) loss.backward() optimizer.step() + # Every 10th epoch, we log the current status. if epoch % 10 == 0: - message = f"epoch={epoch} train_loss={loss.item():.4f}" - if val_input is not None and val_target is not None: - classifier.eval() - with torch.no_grad(): - val_output = classifier(val_input) - val_loss = criterion(val_output, val_target) - val_accuracy = (val_output.argmax(dim=1) == val_target).float().mean().item() - classifier.train() - message += f" val_loss={val_loss.item():.4f} val_acc={val_accuracy:.3f}" - logger.info(message) + train_accuracy = (output.argmax(dim=1) == train_target).float().mean().item() + net.eval() + with torch.no_grad(): + val_output = net(val_input) + val_accuracy = (val_output.argmax(dim=1) == val_target).float().mean().item() -def save_network(output_file: Path, network: Module, label_mapping: list) -> None: - logger.info(f"Saving network {network} to {output_file}") - torch.save({"model_state_dict": network.state_dict(), "label_mapping": label_mapping}, output_file) + logger.info( + f"epoch={epoch} loss={loss.item():.4f} train_acc={train_accuracy:.2f} val_acc={val_accuracy:.2f}" + ) + + logger.info(f"last output={output}") + + # Log the confusion matrix - so we get some better insights into how the training went. + net.eval() + with torch.no_grad(): + predictions = net(val_input).argmax(dim=1) + pred_counts = Counter(predictions.tolist()) + logger.info(f"prediction distribution: {pred_counts}") + for true_class in sorted(set(val_target.tolist())): + mask = val_target == true_class + class_predictions = predictions[mask] + class_counts = Counter(class_predictions.tolist()) + logger.info(f"true={true_class} (n={mask.sum().item()}) -> predicted: {class_counts}")