41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
from pathlib import Path
|
|
|
|
import torch
|
|
from loguru import logger
|
|
from torch import Tensor
|
|
from torch.nn import Linear
|
|
from torch.nn import Module
|
|
from torch.nn import ReLU
|
|
from torch.nn import Sequential
|
|
|
|
|
|
class SleepEpochClassifier(Module):
|
|
"""Per-epoch sleep stage classifier.
|
|
|
|
Input: a fixed-size feature vector summarizing one epoch's context window
|
|
(e.g. motion mean/max, HR mean/variance, time-of-night, previous
|
|
predicted label — whatever features you settle on).
|
|
Output: logits over sleep stage classes for that single epoch.
|
|
|
|
No recurrence/attention — sequence smoothing (majority vote / HMM)
|
|
happens as a separate post-processing step, not inside this network.
|
|
"""
|
|
|
|
def __init__(self, n_features: int, n_classes: int, hidden_size: int = 16):
|
|
super().__init__()
|
|
self.net = Sequential(
|
|
Linear(n_features, hidden_size),
|
|
ReLU(),
|
|
Linear(hidden_size, n_classes),
|
|
)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
# x shape: (batch, n_features) -> returns (batch, n_classes) logits
|
|
return self.net(x)
|
|
|
|
def load(self, input_file: Path) -> list | None:
|
|
payload = torch.load(input_file, weights_only=True)
|
|
state_dict = payload.get("model_state_dict", payload)
|
|
self.load_state_dict(state_dict)
|
|
return payload.get("label_mapping")
|