From 130b97905cacbe48c1381375a5d1993e24718b89 Mon Sep 17 00:00:00 2001 From: Asger Gitz-Johansen Date: Sun, 2 Aug 2026 18:05:26 +0200 Subject: [PATCH] cleanup: move code around --- src/sleep_detection/main.py | 90 ++++++++++++++++++++++++------------- 1 file changed, 60 insertions(+), 30 deletions(-) diff --git a/src/sleep_detection/main.py b/src/sleep_detection/main.py index b08204d..1ee0544 100644 --- a/src/sleep_detection/main.py +++ b/src/sleep_detection/main.py @@ -23,19 +23,39 @@ from sleep_detection.sleep_epoch_classifer import SleepEpochClassifier def main() -> None: """The Main entrypoint function.""" cli = parse_args() - session = load_session(cli.a, cli.H, cli.l) + sessions = load_sessions(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}") + features_per_session, targets_per_session = [], [] + for session in sessions: + inputs, targets = featureize(session) + features_per_session.append(inputs) + targets_per_session.append(targets) - targets = torch.tensor(session["stage"].to_numpy(), dtype=torch.long) - logger.info(f"targets={targets}") + n_sessions = len(features_per_session) + n_train = max(1, int(n_sessions * 0.8)) + perm = torch.randperm(n_sessions) + train_idx, val_idx = perm[:n_train], perm[n_train:] + + train_input = torch.cat([features_per_session[i] for i in train_idx]) + train_target = torch.cat([targets_per_session[i] for i in train_idx]) + if len(val_idx) == 0: + val_input, val_target = train_input, train_target + else: + val_input = torch.cat([features_per_session[i] for i in val_idx]) + val_target = torch.cat([targets_per_session[i] for i in val_idx]) + + # Normalize using train stats only. + mean = train_input.mean(dim=0, keepdim=True) + std = train_input.std(dim=0, keepdim=True) + 1e-8 + train_input = (train_input - mean) / std + val_input = (val_input - mean) / std + + class_counts = torch.bincount(train_target, minlength=3) + class_weights = 1.0 / torch.sqrt(class_counts.float()) + class_weights = class_weights / class_weights.sum() net = SleepEpochClassifier(n_features=3, hidden_size=16, n_classes=3) - train(normalized_inputs, targets, net) + train(train_input, train_target, val_input, val_target, net, class_weights) torch.save(net.state_dict(), cli.output) @@ -44,13 +64,25 @@ def parse_args() -> Namespace: 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("-a", type=Path, nargs="+", required=True, help="Acceleration data file(s).") + parser.add_argument("-H", type=Path, nargs="+", required=True, help="Heartrate data file(s).") + parser.add_argument("-l", type=Path, nargs="+", 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_sessions( + acceleration_files: list[Path], heartrate_files: list[Path], label_files: list[Path] +) -> list[pandas.DataFrame]: + if not (len(acceleration_files) == len(heartrate_files) == len(label_files)): + raise RuntimeError("unmatched session file sets.") + result = [] + for i in range(len(acceleration_files)): + result.append(load_session(acceleration_files[i], heartrate_files[i], label_files[i])) + logger.info(f"sessions={len(result)}") + return result + + def load_session(acceleration_file: Path, heartrate_file: Path, label_file: Path) -> pandas.DataFrame: """Load a session triplet.""" adf = load_acceleration_data(acceleration_file) @@ -71,6 +103,13 @@ def load_session(acceleration_file: Path, heartrate_file: Path, label_file: Path return hdf +def featureize(session: pandas.DataFrame) -> tuple[Tensor, Tensor]: + """Convert one session dataframe into model inputs and targets.""" + inputs = torch.tensor(session[["timestamp", "heartrate", "motion-avg"]].to_numpy(), dtype=torch.float32) + targets = torch.tensor(session["stage"].to_numpy(), dtype=torch.long) + return inputs, targets + + def load_acceleration_data(data_file: Path) -> pandas.DataFrame: """Load acceleration data file.""" logger.info(f"Loading acceleration data file {data_file}") @@ -108,44 +147,35 @@ def normalize_sleep_stage(stage) -> int: if pandas.isna(stage): return int(numpy.nan) value = int(stage) - if value == 0: + if value <= 0: return 0 if 1 <= value <= 4: return 1 if 5 <= value <= 6: return 2 - return int(numpy.nan) + raise ValueError(stage) -def train(input: Tensor, target: Tensor, net: Module) -> None: +def train( + input: Tensor, target: Tensor, val_input: Tensor, val_target: Tensor, net: Module, class_weights: Tensor +) -> 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): + for epoch in range(1000): optimizer.zero_grad() - output = net(train_input) - loss = criterion(output, train_target) + output = net(input) + loss = criterion(output, 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() + train_accuracy = (output.argmax(dim=1) == target).float().mean().item() net.eval() with torch.no_grad():