172 lines
6.5 KiB
Python
172 lines
6.5 KiB
Python
"""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}")
|