cleanup: move code around

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():