wip

src/sleep_detection/main.py

@@ -25,17 +25,37 @@ def main() -> None:
     cli = parse_args()
     sessions = load_sessions(cli.a, cli.H, cli.l)
 
-    inputs = torch.tensor(sessions[["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(sessions["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)
 
 
@@ -53,13 +73,13 @@ def parse_args() -> Namespace:
 
 def load_sessions(
     acceleration_files: list[Path], heartrate_files: list[Path], label_files: list[Path]
-) -> pandas.DataFrame:
-    if len(acceleration_files) != len(heartrate_files) != len(label_files):
+) -> list[pandas.DataFrame]:
+    if not (len(acceleration_files) == len(heartrate_files) == len(label_files)):
         raise RuntimeError("unmatched session file sets.")
-    result = pandas.DataFrame()
+    result = []
     for i in range(len(acceleration_files)):
-        result = pandas.concat([result, load_session(acceleration_files[i], heartrate_files[i], label_files[i])])
-    logger.info(f"sessions={result}")
+        result.append(load_session(acceleration_files[i], heartrate_files[i], label_files[i]))
+    logger.info(f"sessions={len(result)}")
     return result
 
 
@@ -83,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}")
@@ -129,35 +156,26 @@ def normalize_sleep_stage(stage) -> int:
     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():