feat: much better architecture

I understand at least some of this.

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")
+    session = load_session(cli.a, cli.H, cli.l)
-    load_session(cli.a, cli.H, cli.l)
+
-    raise ValueError(f"Unknown command: {cli.command}")
+    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):
+def normalize_sleep_stage(stage) -> int:
-    classifier.train()
+    """Map raw sleep stages to the 3-class target space."""
-    optimizer = SGD(classifier.parameters(), lr=0.01)
+    if pandas.isna(stage):
-    criterion = CrossEntropyLoss()
+        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]
+    train_idx, val_idx = permutation[:n_train], permutation[n_train:]
-    val_idx = 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]
+    # Tweak the loss function.
-    train_target = target[train_idx]
+    class_counts = torch.bincount(train_target)
-    val_input = input[val_idx] if len(val_idx) else None
+    class_weights = 1.0 / torch.sqrt(class_counts.float())
-    val_target = target[val_idx] if len(val_idx) else None
+    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}"
+            train_accuracy = (output.argmax(dim=1) == train_target).float().mean().item()
-            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)
 
+            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(
-    logger.info(f"Saving network {network} to {output_file}")
+                f"epoch={epoch} loss={loss.item():.4f} train_acc={train_accuracy:.2f} val_acc={val_accuracy:.2f}"
-    torch.save({"model_state_dict": network.state_dict(), "label_mapping": label_mapping}, output_file)
+            )
+
+    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}")