feat: first training attempt
It still sucks. I'm sure the the feature embedding is terribly wrong.
This commit is contained in:
+340
-2
@@ -1,6 +1,344 @@
|
||||
"""Main Entrypoint codebase."""
|
||||
|
||||
import sys
|
||||
from argparse import ArgumentParser
|
||||
from argparse import Namespace
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
def main():
|
||||
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
|
||||
|
||||
SLEEP_STAGE_LABELS = [0, 1, 2]
|
||||
|
||||
|
||||
def normalize_sleep_stage(stage) -> float:
|
||||
"""Map raw sleep stages to the 3-class target space.
|
||||
|
||||
Raw labels:
|
||||
0 -> 0
|
||||
1, 2, 3, 4 -> 1
|
||||
5, 6 -> 2
|
||||
Other values are treated as unknown.
|
||||
"""
|
||||
if pandas.isna(stage):
|
||||
return numpy.nan
|
||||
|
||||
value = int(stage)
|
||||
if value == 0:
|
||||
return 0.0
|
||||
if 1 <= value <= 4:
|
||||
return 1.0
|
||||
if 5 <= value <= 6:
|
||||
return 2.0
|
||||
|
||||
return numpy.nan
|
||||
|
||||
|
||||
@logger.catch
|
||||
def main() -> None:
|
||||
"""The Main entrypoint function."""
|
||||
print("Hello from sleep-detection!")
|
||||
cli = parse_args()
|
||||
logger.info("welcome")
|
||||
if cli.command == "train":
|
||||
sessions = load_sessions(cli.a, cli.H, cli.l)
|
||||
epochs, features, target, label_mapping = featureize_sessions(sessions)
|
||||
|
||||
c = SleepEpochClassifier(
|
||||
n_features=features.shape[1],
|
||||
n_classes=len(SLEEP_STAGE_LABELS),
|
||||
hidden_size=8,
|
||||
)
|
||||
if cli.input is not None:
|
||||
c.load(cli.input)
|
||||
|
||||
logger.info(f"Using a {c.__class__.__name__}")
|
||||
|
||||
train(features, target, c)
|
||||
save_network(cli.output, c, label_mapping)
|
||||
return
|
||||
|
||||
if cli.command == "execute":
|
||||
adf = load_acceleration_data(cli.a)
|
||||
hdf = load_heartrate_data(cli.H)
|
||||
ldf = load_labels_data(cli.l)
|
||||
|
||||
df = adf.join(hdf, how="outer").join(ldf, how="outer").sort_index()
|
||||
epochs, features, target, label_mapping = featureize(df)
|
||||
|
||||
c = SleepEpochClassifier(
|
||||
n_features=features.shape[1],
|
||||
n_classes=len(SLEEP_STAGE_LABELS),
|
||||
hidden_size=8,
|
||||
)
|
||||
loaded_label_mapping = c.load(cli.model)
|
||||
logger.info(f"Using a {c.__class__.__name__}")
|
||||
execute(epochs, features, target, c, loaded_label_mapping or label_mapping)
|
||||
return
|
||||
|
||||
raise ValueError(f"Unknown command: {cli.command}")
|
||||
|
||||
|
||||
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")
|
||||
command = parser.add_subparsers(dest="command")
|
||||
|
||||
train = command.add_parser("train", help="Train the model.")
|
||||
train.add_argument("-a", type=Path, nargs="+", required=True, help="Acceleration data file(s).")
|
||||
train.add_argument("-H", type=Path, nargs="+", required=True, help="Heartrate data file(s).")
|
||||
train.add_argument("-l", type=Path, nargs="+", required=True, help="Labels data file(s).")
|
||||
train.add_argument("-o", "--output", type=Path, default=None, required=True, help="Output .pt file.")
|
||||
train.add_argument("-i", "--input", type=Path, default=None, required=False, help="Input .pt file.")
|
||||
|
||||
execute = command.add_parser("execute", help="Execute the model.")
|
||||
execute.add_argument("-m", "--model", type=Path, default=None, required=True, help="Input .pr file to use.")
|
||||
execute.add_argument("-a", type=Path, default=None, required=True, help="Acceleration data file.")
|
||||
execute.add_argument("-H", type=Path, default=None, required=True, help="Heartrate data file.")
|
||||
execute.add_argument("-l", type=Path, default=None, required=True, help="Labels data file.")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def featureize(data: pandas.DataFrame) -> tuple[pandas.DataFrame, Tensor, Tensor, list]:
|
||||
"""Convert loaded data into input and target tensors.
|
||||
|
||||
Args:
|
||||
data: Loaded data, should contain the keys: [timestamp, heartrate, motion, x, y, z, stage]
|
||||
|
||||
"""
|
||||
logger.info(f"Featurizing the data...")
|
||||
epochs = build_epochs(data)
|
||||
|
||||
epochs["hr_std"] = epochs["hr_std"].fillna(0.0)
|
||||
epochs = epochs.dropna(subset=["hr_mean", "motion_mean", "motion_max", "stage"])
|
||||
|
||||
epochs["seconds_since_start"] = epochs.index.total_seconds()
|
||||
|
||||
CONTEXT = 2 # epochs each side -> 5 epochs total = 25 min window
|
||||
FEATURE_COLS = ["hr_mean", "hr_std", "motion_mean", "motion_max", "seconds_since_start"]
|
||||
|
||||
features = epochs[FEATURE_COLS].to_numpy(dtype=numpy.float32)
|
||||
labels = epochs["stage"].astype(numpy.int64).to_numpy()
|
||||
|
||||
X, y = [], []
|
||||
for i in range(CONTEXT, len(epochs) - CONTEXT):
|
||||
window = features[i - CONTEXT : i + CONTEXT + 1] # shape: (2*CONTEXT+1, n_feats)
|
||||
X.append(window.flatten()) # -> 1D vector
|
||||
y.append(labels[i])
|
||||
|
||||
if not X:
|
||||
raise ValueError("Not enough complete epochs to build a training set.")
|
||||
|
||||
X = torch.tensor(numpy.stack(X), dtype=torch.float32) # shape: (n_epochs, (2*CONTEXT+1)*n_feats)
|
||||
y = torch.tensor(y, dtype=torch.long) # shape: (n_epochs,)
|
||||
return epochs, X, y, SLEEP_STAGE_LABELS
|
||||
|
||||
|
||||
def featureize_sessions(dataframes: Sequence[pandas.DataFrame]) -> tuple[pandas.DataFrame, Tensor, Tensor, list]:
|
||||
"""Convert multiple loaded sessions into input and target tensors."""
|
||||
epochs_per_session = []
|
||||
for data in dataframes:
|
||||
epochs = build_epochs(data)
|
||||
|
||||
epochs["hr_std"] = epochs["hr_std"].fillna(0.0)
|
||||
epochs = epochs.dropna(subset=["hr_mean", "motion_mean", "motion_max", "stage"])
|
||||
epochs["seconds_since_start"] = epochs.index.total_seconds()
|
||||
epochs_per_session.append(epochs)
|
||||
|
||||
if not epochs_per_session:
|
||||
raise ValueError("No sessions provided.")
|
||||
|
||||
combined_epochs = pandas.concat(epochs_per_session)
|
||||
|
||||
CONTEXT = 2 # epochs each side -> 5 epochs total = 25 min window
|
||||
FEATURE_COLS = ["hr_mean", "hr_std", "motion_mean", "motion_max", "seconds_since_start"]
|
||||
|
||||
X, y = [], []
|
||||
for epochs in epochs_per_session:
|
||||
features = epochs[FEATURE_COLS].to_numpy(dtype=numpy.float32)
|
||||
labels = epochs["stage"].astype(numpy.int64).to_numpy()
|
||||
|
||||
for i in range(CONTEXT, len(epochs) - CONTEXT):
|
||||
window = features[i - CONTEXT : i + CONTEXT + 1] # shape: (2*CONTEXT+1, n_feats)
|
||||
X.append(window.flatten()) # -> 1D vector
|
||||
y.append(labels[i])
|
||||
|
||||
if not X:
|
||||
raise ValueError("Not enough complete epochs to build a training set.")
|
||||
|
||||
X = torch.tensor(numpy.stack(X), dtype=torch.float32) # shape: (n_epochs, (2*CONTEXT+1)*n_feats)
|
||||
y = torch.tensor(y, dtype=torch.long) # shape: (n_epochs,)
|
||||
return combined_epochs, X, y, SLEEP_STAGE_LABELS
|
||||
|
||||
|
||||
def load_sessions(
|
||||
acceleration_files: Sequence[Path], heartrate_files: Sequence[Path], label_files: Sequence[Path]
|
||||
) -> list[pandas.DataFrame]:
|
||||
"""Load multiple session triplets."""
|
||||
if not (len(acceleration_files) == len(heartrate_files) == len(label_files)):
|
||||
raise ValueError("Acceleration, heartrate, and label file counts must match.")
|
||||
|
||||
sessions = []
|
||||
for acceleration_file, heartrate_file, label_file in zip(acceleration_files, heartrate_files, label_files):
|
||||
adf = load_acceleration_data(acceleration_file)
|
||||
hdf = load_heartrate_data(heartrate_file)
|
||||
ldf = load_labels_data(label_file)
|
||||
sessions.append(adf.join(hdf, how="outer").join(ldf, how="outer").sort_index())
|
||||
|
||||
return sessions
|
||||
|
||||
|
||||
def build_epochs(data: pandas.DataFrame) -> pandas.DataFrame:
|
||||
"""Aggregate raw signals into 5-minute epochs."""
|
||||
EPOCH_LEN = pandas.Timedelta(minutes=5)
|
||||
|
||||
return data.resample(EPOCH_LEN).agg(
|
||||
hr_mean=("hr", "mean"),
|
||||
hr_std=("hr", "std"),
|
||||
motion_mean=("motion", "mean"),
|
||||
motion_max=("motion", "max"),
|
||||
stage=("stage", lambda s: s.dropna().mode().iloc[0] if len(s.dropna()) else numpy.nan),
|
||||
)
|
||||
|
||||
|
||||
def execute(epochs: pandas.DataFrame, input: Tensor, target: Tensor, classifier: Module, label_mapping: list) -> None:
|
||||
"""Run the model and emit per-epoch predictions."""
|
||||
classifier.eval()
|
||||
with torch.no_grad():
|
||||
output = classifier(input)
|
||||
predicted = output.argmax(dim=1)
|
||||
|
||||
decoded_predicted = [label_mapping[index] for index in predicted.tolist()]
|
||||
decoded_target = [label_mapping[index] for index in target.tolist()]
|
||||
|
||||
report = pandas.DataFrame(
|
||||
{
|
||||
"timestamp": epochs.index[2 : len(epochs) - 2].astype("timedelta64[s]") / numpy.timedelta64(1, "s"),
|
||||
"predicted": decoded_predicted,
|
||||
"actual": decoded_target,
|
||||
}
|
||||
)
|
||||
report.to_csv(sys.stdout, index=False)
|
||||
|
||||
accuracy = (predicted == target).float().mean().item()
|
||||
logger.info(f"accuracy={accuracy:.3f}")
|
||||
|
||||
|
||||
def load_acceleration_data(data_file: Path) -> pandas.DataFrame:
|
||||
"""Load acceleration data file.
|
||||
|
||||
Expected input format: <timestamp> <x> <y> <z>
|
||||
|
||||
Args:
|
||||
data_file: The filepath for the file containing acceleration data.
|
||||
|
||||
Returns:
|
||||
A pandas data frame with the loaded acceleration data.
|
||||
|
||||
"""
|
||||
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["timestamp"] = pandas.to_numeric(df["timestamp"], errors="coerce")
|
||||
df = df.dropna(subset=["timestamp"])
|
||||
df = df.sort_values("timestamp").set_index(pandas.to_timedelta(df["timestamp"], unit="s"))
|
||||
df = df.drop(columns=["timestamp"])
|
||||
df.index.name = "timestamp"
|
||||
# df["motion"] = df[["x", "y", "z"]].pow(2).sum(axis=1).pow(0.5)
|
||||
df["motion"] = df[["x", "y", "z"]].diff().pow(2).sum(axis=1).pow(0.5)
|
||||
return df
|
||||
|
||||
|
||||
def load_heartrate_data(data_file: Path) -> pandas.DataFrame:
|
||||
"""Load heart rate data file.
|
||||
|
||||
Expected input format: <timestamp>,<heartrate>
|
||||
|
||||
Args:
|
||||
data_file: The filepath for the file containing heartrate data.
|
||||
|
||||
Returns:
|
||||
A pandas data frame with the loaded heartrate data.
|
||||
|
||||
"""
|
||||
logger.info(f"Loading heartrate date file {data_file}")
|
||||
df = pandas.read_csv(data_file, sep=",", names=["timestamp", "hr"], header=None)
|
||||
df["timestamp"] = pandas.to_numeric(df["timestamp"], errors="coerce")
|
||||
df = df.dropna(subset=["timestamp"])
|
||||
df = df.sort_values("timestamp").set_index(pandas.to_timedelta(df["timestamp"], unit="s"))
|
||||
df = df.drop(columns=["timestamp"])
|
||||
df.index.name = "timestamp"
|
||||
return df
|
||||
|
||||
|
||||
def load_labels_data(data_file: Path) -> pandas.DataFrame:
|
||||
"""Load label data file.
|
||||
|
||||
Expected input format: <timestamp>,<label>
|
||||
|
||||
Args:
|
||||
data_file: The filepath for the file containing label data.
|
||||
|
||||
Returns:
|
||||
A pandas data frame with the loaded label data.
|
||||
|
||||
"""
|
||||
logger.info(f"Loading label date file {data_file}")
|
||||
df = pandas.read_csv(data_file, sep=r"\s+", names=["timestamp", "stage"], header=None)
|
||||
df["timestamp"] = pandas.to_numeric(df["timestamp"], errors="coerce")
|
||||
df["stage"] = pandas.to_numeric(df["stage"], errors="coerce").map(normalize_sleep_stage)
|
||||
df = df.dropna(subset=["timestamp"])
|
||||
df = df.sort_values("timestamp").set_index(pandas.to_timedelta(df["timestamp"], unit="s"))
|
||||
df = df.drop(columns=["timestamp"])
|
||||
df.index.name = "timestamp"
|
||||
return df
|
||||
|
||||
|
||||
def train(input: Tensor, target: Tensor, classifier: Module):
|
||||
classifier.train()
|
||||
optimizer = SGD(classifier.parameters(), lr=0.01)
|
||||
criterion = CrossEntropyLoss()
|
||||
|
||||
n_samples = input.shape[0]
|
||||
n_train = max(1, int(n_samples * 0.8))
|
||||
permutation = torch.randperm(n_samples)
|
||||
train_idx = permutation[:n_train]
|
||||
val_idx = permutation[n_train:]
|
||||
|
||||
train_input = input[train_idx]
|
||||
train_target = target[train_idx]
|
||||
val_input = input[val_idx] if len(val_idx) else None
|
||||
val_target = target[val_idx] if len(val_idx) else None
|
||||
|
||||
for epoch in range(100):
|
||||
optimizer.zero_grad()
|
||||
output = classifier(train_input)
|
||||
loss = criterion(output, train_target)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
if epoch % 10 == 0:
|
||||
message = f"epoch={epoch} train_loss={loss.item():.4f}"
|
||||
if val_input is not None and val_target is not None:
|
||||
classifier.eval()
|
||||
with torch.no_grad():
|
||||
val_output = classifier(val_input)
|
||||
val_loss = criterion(val_output, val_target)
|
||||
val_accuracy = (val_output.argmax(dim=1) == val_target).float().mean().item()
|
||||
classifier.train()
|
||||
message += f" val_loss={val_loss.item():.4f} val_acc={val_accuracy:.3f}"
|
||||
logger.info(message)
|
||||
|
||||
|
||||
def save_network(output_file: Path, network: Module, label_mapping: list) -> None:
|
||||
logger.info(f"Saving network {network} to {output_file}")
|
||||
torch.save({"model_state_dict": network.state_dict(), "label_mapping": label_mapping}, output_file)
|
||||
|
||||
Reference in New Issue
Block a user