"""Main Entrypoint codebase.""" import sys from argparse import ArgumentParser from argparse import Namespace from collections import Counter from collections.abc import Sequence from pathlib import Path import numpy import pandas import torch from loguru import logger from torch import Tensor from torch.nn import CrossEntropyLoss from torch.nn import Module from torch.optim import SGD from sleep_detection.sleep_epoch_classifer import SleepEpochClassifier @logger.catch def main() -> None: """The Main entrypoint function.""" cli = parse_args() 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 load_session(acceleration_file: Path, heartrate_file: Path, label_file: Path) -> pandas.DataFrame: """Load a session triplet.""" adf = load_acceleration_data(acceleration_file) hdf = load_heartrate_data(heartrate_file) ldf = load_labels_data(label_file) # Average the motion into the sparser heartrate dataset. bin_edges = [-float("inf")] + hdf.index.tolist() logger.info(f"({len(hdf)})") adf["bin"] = pandas.cut(adf.index, bins=bin_edges, labels=hdf.index, right=True) hdf["motion-avg"] = adf.groupby("bin", observed=True)["motion"].mean() hdf = hdf.dropna() # Label the heartrate dataset. ldf.index = ldf.index.astype(float) hdf = pandas.merge_asof(hdf, ldf, on="timestamp", direction="backward") logger.info(f"\n{hdf}") return hdf def load_acceleration_data(data_file: Path) -> pandas.DataFrame: """Load acceleration data file.""" logger.info(f"Loading acceleration data file {data_file}") df = pandas.read_csv(data_file, sep=r"\s+", names=["timestamp", "x", "y", "z"], header=None) df = df.drop(df[df["timestamp"] < 0].index).drop_duplicates() df["motion"] = df[["x", "y", "z"]].diff().pow(2).sum(axis=1).pow(0.5) df = df.set_index("timestamp") logger.info(f"({len(df)})\n{df}") return df def load_heartrate_data(data_file: Path) -> pandas.DataFrame: """Load heart rate data file.""" logger.info(f"Loading heartrate date file {data_file}") df = pandas.read_csv(data_file, sep=",", names=["timestamp", "heartrate"], header=None) 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 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 = df.drop(df[df["timestamp"] < 0].index).drop_duplicates() df = df.set_index("timestamp") logger.info(f"({len(df)})\n{df}") return df 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, 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] # 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 = net(train_input) loss = criterion(output, train_target) loss.backward() optimizer.step() # Every 10th epoch, we log the current status. if epoch % 10 == 0: 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() 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}")