diff --git a/configs/model/tcn_single_joint.yaml b/configs/model/tcn_single_joint.yaml new file mode 100644 index 0000000..9ea0aac --- /dev/null +++ b/configs/model/tcn_single_joint.yaml @@ -0,0 +1,39 @@ +# TCN Single Joint Model Configuration +# Scaled version of default tcn configuration + +name: tcn_single_joint +type: tcn + +# Architecture parameters +architecture: + input_size: 7 # 6 IMU features + 1 joint angle + output_size: 1 # 1 joint moment for testing + + # Single joint parameters + num_channels: [6, 6, 6, 6, 6] # 5 layers with 6 channels each -> this is scaled from the default 25 since we are dealing with 1/4 of the information (reduce overfitting) + kernel_size: 7 + dropout: 0.2 + spatial_dropout: false + activation: ReLU + norm: weight_norm + +# Same Computed properties +effective_history: 187 # timesteps +receptive_field_ms: 1870 # milliseconds at 100Hz + +# Same Normalization +normalization: + enabled: false + use_dataset_stats: true + +# Same weight initialization +initialization: + conv_std: 0.01 + linear_std: 0.01 + +data_mapping: + joint_index: 0 # UPDATE to reflect joint being tested. Hip L = 0, Hip R = 1, Knee L = 2, Knee R = 3 + joint_features: [0, 1, 2, 3, 4, 5, 6] # UPDATE to reflect columns containing features of the joint being tested + + +estimated_params: 47000 diff --git a/configs/pilots/default.json b/configs/pilots/default.json new file mode 100644 index 0000000..90f3382 --- /dev/null +++ b/configs/pilots/default.json @@ -0,0 +1,4 @@ +{ + "name": "default", + "mass_kg": 1.0 +} diff --git a/deploy/exoskeleton-server.service b/deploy/exoskeleton-server.service index 040ee0a..489727d 100644 --- a/deploy/exoskeleton-server.service +++ b/deploy/exoskeleton-server.service @@ -7,6 +7,7 @@ ExecStart=/home/dylan-exo/ml_stuff/exoskeleton-ai/.venv/bin/uvicorn scripts.serv WorkingDirectory=/home/dylan-exo/ml_stuff/exoskeleton-ai Restart=on-failure User=dylan-exo +Environment=PILOT=default [Install] WantedBy=multi-user.target diff --git a/scripts/export_onnx.py b/scripts/export_onnx.py index f4904d8..9ee07f0 100644 --- a/scripts/export_onnx.py +++ b/scripts/export_onnx.py @@ -21,7 +21,7 @@ DEFAULT_MODEL_PATH = "outputs/2026-01-08/18-29-05/best_model.pt" -def load_model(model_path: Path) -> tuple[torch.nn.Module, int]: +def load_model(model_path: Path): checkpoint = torch.load(model_path, map_location="cpu") config_path = model_path.parent / "config.yaml" @@ -33,22 +33,24 @@ def load_model(model_path: Path) -> tuple[torch.nn.Module, int]: model.load_state_dict(checkpoint["model_state_dict"]) model.eval() - window_size: int = config.model.effective_history - return model, window_size + return model, config def export(model_path: Path, output_path: Path) -> None: print(f"Loading model: {model_path}") - model, window_size = load_model(model_path) + model, config = load_model(model_path) + window_size: int = config.model.effective_history num_params = sum(p.numel() for p in model.parameters()) + input_size: int = config.model.architecture.input_size + output_size: int = config.model.architecture.output_size print(f" Parameters: {num_params:,}") print(f" Window size: {window_size} timesteps") - print(f" Input shape: (1, {window_size}, 28)") - print(f" Output shape: (1, {window_size}, 4)") + print(f" Input shape: (1, {window_size}, {input_size})") + print(f" Output shape: (1, {window_size}, {output_size})") # Representative input for tracing - dummy_input = torch.randn(1, window_size, 28) + dummy_input = torch.randn(1, window_size, input_size) with torch.no_grad(): pytorch_output = model(dummy_input) diff --git a/scripts/predict_msgpack.cpp b/scripts/predict_msgpack.cpp new file mode 100644 index 0000000..b5d5728 --- /dev/null +++ b/scripts/predict_msgpack.cpp @@ -0,0 +1,82 @@ +/** + * predict_msgpack.cpp + * + * Minimal C++ client that sends a single msgpack inference request + * to the exoskeleton TCN server and prints the predicted torques. + * + * Compile: + * g++ -O2 -std=c++17 predict_msgpack.cpp -lcurl -o predict_msgpack + * + * Run: + * ./predict_msgpack + */ + +#include +#include +#include + +#include +#include + +static const char* MSGPACK_URL = "http://127.0.0.1:8000/predict_msgpack?model=single_joint"; +static const int WINDOW_SIZE = 187; +static const int FEATURE_COUNT = 7; // hip_left only: 6 IMU + 1 joint angle + +static size_t write_callback(char* ptr, size_t size, size_t nmemb, void* userdata) { + auto* buf = static_cast(userdata); + buf->append(ptr, size * nmemb); + return size * nmemb; +} + +int main() { + // Build flat input: WINDOW_SIZE * FEATURE_COUNT float32 values (zeroed) + std::vector flat(WINDOW_SIZE * FEATURE_COUNT, 0.0f); + + // Pack with msgpack + msgpack::sbuffer buf; + msgpack::pack(buf, flat); + + // Send POST request + curl_global_init(CURL_GLOBAL_DEFAULT); + CURL* curl = curl_easy_init(); + if (!curl) { + std::cerr << "Failed to init curl\n"; + return 1; + } + + std::string response_body; + struct curl_slist* headers = nullptr; + headers = curl_slist_append(headers, "Content-Type: application/x-msgpack"); + + curl_easy_setopt(curl, CURLOPT_URL, MSGPACK_URL); + curl_easy_setopt(curl, CURLOPT_POST, 1L); + curl_easy_setopt(curl, CURLOPT_POSTFIELDS, buf.data()); + curl_easy_setopt(curl, CURLOPT_POSTFIELDSIZE, (long)buf.size()); + curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers); + curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_callback); + curl_easy_setopt(curl, CURLOPT_WRITEDATA, &response_body); + curl_easy_setopt(curl, CURLOPT_TIMEOUT, 5L); + + CURLcode res = curl_easy_perform(curl); + long http_code = 0; + curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &http_code); + + curl_slist_free_all(headers); + curl_easy_cleanup(curl); + curl_global_cleanup(); + + if (res != CURLE_OK || http_code != 200) { + std::cerr << "Request failed (HTTP " << http_code << "): " << curl_easy_strerror(res) << "\n"; + return 1; + } + + // Unpack response map + msgpack::object_handle oh = msgpack::unpack(response_body.data(), response_body.size()); + std::map result; + oh.get().convert(result); + + std::cout << "hip_left: " << result["hip_left"] << " Nm\n"; + std::cout << "inference_ms: " << result["inference_ms"] << " ms\n"; + + return 0; +} diff --git a/scripts/server.py b/scripts/server.py index f2f75e2..352769e 100644 --- a/scripts/server.py +++ b/scripts/server.py @@ -2,25 +2,30 @@ FastAPI inference server for the exoskeleton TCN model. Runs on the Raspberry Pi. The C++ control system sends a POST request to -/predict with a full (187, 28) window of sensor data and receives 4 motor -torque predictions in response. +/predict with a full sensor window and receives motor torque predictions +(in Nm) in response. Usage: uvicorn scripts.server:app --host 0.0.0.0 --port 8000 uvicorn scripts.server:app --host 127.0.0.1 --port 8000 # localhost only +Select model via query param (default: full): + POST /predict?model=full + POST /predict?model=single_joint + Dependencies (not in pyproject.toml, install on Pi): pip install fastapi uvicorn[standard] onnxruntime numpy msgpack -Input shape: (1, WINDOW_SIZE, 28) — WINDOW_SIZE defaults to 187 for standard TCN -Output shape: (1, WINDOW_SIZE, 4) — server returns only the final timestep's 4 values +Input shape: (1, window_size, feature_count) — per model registry entry +Output shape: (1, window_size, output_count) — server returns only the final timestep's values, + denormalized to physical torque units (Nm) """ +import json +import os +import time from contextlib import asynccontextmanager from pathlib import Path -from typing import Annotated - -import time import msgpack import numpy as np @@ -32,41 +37,91 @@ # Configuration # --------------------------------------------------------------------------- -MODEL_PATH = "outputs/model.onnx" - -# Must match the effective receptive field of the deployed model variant: -# standard TCN (kernel=7, 5 layers): 187 -# small TCN (kernel=5, 4 layers): 61 -# large TCN (kernel=9, 6 layers): 505 -WINDOW_SIZE = 187 -FEATURE_COUNT = 28 # 24 IMU + 4 joint angles -OUTPUT_COUNT = 4 # hip_l, hip_r, knee_l, knee_r (Nm/kg) +JOINT_NAMES = ["hip_left", "hip_right", "knee_left", "knee_right"] + +NORM_STATS_PATH = "data/processed/phase1/normalization_stats.json" +PILOTS_DIR = Path("configs/pilots") + +MODEL_REGISTRY: dict[str, dict] = { + "full": { + "onnx_path": "outputs/model.onnx", + "window_size": 187, + "feature_count": 28, + "joint_indices": [0, 1, 2, 3], + }, + "single_joint": { + "onnx_path": "outputs/model_single_joint.onnx", + "window_size": 187, + "feature_count": 7, + "joint_indices": [0], # hip_left only + "feature_indices": [0, 1, 2, 3, 4, 5, 6], # columns in normalization_stats inputs + }, +} +DEFAULT_MODEL = "full" # --------------------------------------------------------------------------- -# App lifespan: load model once at startup +# App lifespan: load models and normalization stats once at startup # --------------------------------------------------------------------------- -_session: ort.InferenceSession | None = None -_input_name: str | None = None +_inputs_mean: np.ndarray | None = None +_inputs_std: np.ndarray | None = None +_targets_mean: np.ndarray | None = None +_targets_std: np.ndarray | None = None +_pilot_mass_kg: float = 1.0 +_sessions: dict[str, ort.InferenceSession] = {} +_input_names: dict[str, str] = {} @asynccontextmanager async def lifespan(app: FastAPI): - global _session, _input_name - - model_path = Path(MODEL_PATH) - if not model_path.exists(): - raise RuntimeError(f"Model file not found: {model_path.resolve()}") - - print(f"Loading ONNX model: {model_path.resolve()}") - _session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) - _input_name = _session.get_inputs()[0].name - print(f"Model loaded. Input: '{_input_name}' | Expected shape: (1, {WINDOW_SIZE}, {FEATURE_COUNT})") + global _inputs_mean, _inputs_std, _targets_mean, _targets_std, _pilot_mass_kg, _sessions, _input_names + + # Load pilot config + pilot_name = os.environ.get("PILOT", "default") + pilot_path = PILOTS_DIR / f"{pilot_name}.json" + if not pilot_path.exists(): + print(f"WARNING: Pilot config '{pilot_path}' not found, falling back to default.json") + pilot_path = PILOTS_DIR / "default.json" + pilot_cfg = json.loads(pilot_path.read_text()) + _pilot_mass_kg = float(pilot_cfg["mass_kg"]) + print(f"Pilot: {pilot_cfg.get('name', pilot_name)} | mass = {_pilot_mass_kg} kg") + + # Load normalization stats + stats_path = Path(NORM_STATS_PATH) + if not stats_path.exists(): + raise RuntimeError(f"Normalization stats not found: {stats_path.resolve()}") + stats = json.loads(stats_path.read_text()) + _inputs_mean = np.array(stats["inputs"]["mean"], dtype=np.float32) + _inputs_std = np.array(stats["inputs"]["std"], dtype=np.float32) + _targets_mean = np.array(stats["targets"]["mean"], dtype=np.float32) + _targets_std = np.array(stats["targets"]["std"], dtype=np.float32) + print(f"Loaded normalization stats from {stats_path.resolve()}") + + # Load all registered model sessions + for name, cfg in MODEL_REGISTRY.items(): + model_path = Path(cfg["onnx_path"]) + if not model_path.exists(): + print(f"WARNING: Model '{name}' not found at {model_path.resolve()}, skipping.") + continue + print(f"Loading ONNX model '{name}': {model_path.resolve()}") + session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + _sessions[name] = session + _input_names[name] = session.get_inputs()[0].name + print( + f" Loaded '{name}': input='{_input_names[name]}' " + f"shape=(1, {cfg['window_size']}, {cfg['feature_count']})" + ) + + if DEFAULT_MODEL not in _sessions: + raise RuntimeError( + f"Default model '{DEFAULT_MODEL}' could not be loaded. " + "Ensure the ONNX file exists before starting the server." + ) yield - _session = None - _input_name = None + _sessions.clear() + _input_names.clear() print("Server shutting down.") @@ -76,100 +131,189 @@ async def lifespan(app: FastAPI): ) # --------------------------------------------------------------------------- -# Request / response schemas +# Inference helper +# --------------------------------------------------------------------------- + + +def _run_inference(model_name: str, window: np.ndarray) -> dict: + """ + Run inference with the named model and return a dict of denormalized predictions. + + Args: + model_name: Key into MODEL_REGISTRY / _sessions. + window: Input array of shape (1, window_size, feature_count). + + Returns: + Dict with joint name keys (Nm values) and 'inference_ms'. + """ + cfg = MODEL_REGISTRY[model_name] + session = _sessions[model_name] + input_name = _input_names[model_name] + + # Normalize inputs: client sends raw sensor values, model expects z-scores + # If the model uses a feature subset, slice the stats to match + feature_indices = cfg.get("feature_indices") + if feature_indices is not None: + inp_mean = _inputs_mean[feature_indices] # type: ignore[index] + inp_std = _inputs_std[feature_indices] # type: ignore[index] + else: + inp_mean = _inputs_mean + inp_std = _inputs_std + window = (window - inp_mean) / inp_std + + t0 = time.perf_counter() + outputs = session.run(None, {input_name: window}) + inference_ms = (time.perf_counter() - t0) * 1000 + + preds_norm = outputs[0][0, -1, :] # shape: (output_count,) + + # Denormalize from z-score → Nm/kg, then scale by pilot mass → Nm + indices = cfg["joint_indices"] + preds_nm_per_kg = preds_norm * _targets_std[indices] + _targets_mean[indices] # type: ignore[index] + preds_phys = preds_nm_per_kg * _pilot_mass_kg + + result: dict = {JOINT_NAMES[j]: float(preds_phys[i]) for i, j in enumerate(indices)} + result["inference_ms"] = round(inference_ms, 3) + return result + + +def _resolve_model(model: str) -> str: + """Validate the model query param and return it, or raise HTTP 400/503.""" + if model not in MODEL_REGISTRY: + raise HTTPException( + status_code=400, + detail=f"Unknown model '{model}'. Available: {list(MODEL_REGISTRY.keys())}", + ) + if model not in _sessions: + raise HTTPException( + status_code=503, + detail=f"Model '{model}' is registered but not loaded (ONNX file missing at startup).", + ) + return model + + +# --------------------------------------------------------------------------- +# Request schema (validated dynamically against the selected model) # --------------------------------------------------------------------------- class PredictRequest(BaseModel): - # List of WINDOW_SIZE rows, each with FEATURE_COUNT values. - # C++ sends: {"window": [[f0..f27], [f0..f27], ...]} (187 rows) - window: Annotated[ - list[list[float]], - Field(description=f"Sensor window: {WINDOW_SIZE} timesteps × {FEATURE_COUNT} features"), - ] + window: list[list[float]] = Field(description="Sensor window: timesteps × features") + + # model_name is injected before validation by the endpoint + model_config = {"arbitrary_types_allowed": True} + + # Stored after endpoint sets it before calling model_validate + _model_name: str = "full" @model_validator(mode="after") def validate_shape(self) -> "PredictRequest": + # Use the model name stored on the instance (set by endpoint before validation) + model_name = getattr(self, "_model_name", DEFAULT_MODEL) + cfg = MODEL_REGISTRY.get(model_name, MODEL_REGISTRY[DEFAULT_MODEL]) + window_size = cfg["window_size"] + feature_count = cfg["feature_count"] + rows = len(self.window) - if rows != WINDOW_SIZE: - raise ValueError(f"Expected {WINDOW_SIZE} timesteps, got {rows}") + if rows != window_size: + raise ValueError(f"Expected {window_size} timesteps, got {rows}") for i, row in enumerate(self.window): - if len(row) != FEATURE_COUNT: + if len(row) != feature_count: raise ValueError( - f"Timestep {i}: expected {FEATURE_COUNT} features, got {len(row)}" + f"Timestep {i}: expected {feature_count} features, got {len(row)}" ) return self -class PredictResponse(BaseModel): - # Torque predictions for the 4 joints at the latest timestep (Nm/kg) - hip_left: float - hip_right: float - knee_left: float - knee_right: float - inference_ms: float # ONNX session.run() time only, excludes HTTP/JSON overhead - - # --------------------------------------------------------------------------- # Endpoints # --------------------------------------------------------------------------- -def _run_inference(window_bytes: np.ndarray) -> tuple[np.ndarray, float]: - """Shared inference logic. Input must already be shaped (1, WINDOW_SIZE, FEATURE_COUNT).""" - t0 = time.perf_counter() - outputs = _session.run(None, {_input_name: window_bytes}) # type: ignore[index] - inference_ms = (time.perf_counter() - t0) * 1000 - preds = outputs[0][0, -1, :] - return preds, inference_ms - - @app.get("/health") -def health() -> dict[str, str]: +def health() -> dict: """Quick liveness check — the C++ side can poll this on startup.""" - if _session is None: + if DEFAULT_MODEL not in _sessions: raise HTTPException(status_code=503, detail="Model not loaded") - return {"status": "ok"} + return { + "status": "ok", + "models_loaded": list(_sessions.keys()), + "pilot_mass_kg": _pilot_mass_kg, + } -@app.post("/predict", response_model=PredictResponse) -def predict(req: PredictRequest) -> PredictResponse: - if _session is None or _input_name is None: - raise HTTPException(status_code=503, detail="Model not loaded") +@app.post("/predict") +def predict(req: PredictRequest, model: str = DEFAULT_MODEL) -> dict: + """ + JSON inference endpoint. + + Query params: + model: Model name from MODEL_REGISTRY (default: 'full') + + Request body: + {"window": [[f0..fN], ...]} — shape (window_size, feature_count) + + Response: + {"hip_left": , ..., "inference_ms": } + """ + model_name = _resolve_model(model) + + cfg = MODEL_REGISTRY[model_name] + # Validate shape against the selected model + rows = len(req.window) + if rows != cfg["window_size"]: + raise HTTPException( + status_code=422, + detail=f"Expected {cfg['window_size']} timesteps for model '{model_name}', got {rows}", + ) + for i, row in enumerate(req.window): + if len(row) != cfg["feature_count"]: + raise HTTPException( + status_code=422, + detail=( + f"Timestep {i}: expected {cfg['feature_count']} features " + f"for model '{model_name}', got {len(row)}" + ), + ) input_array = np.array(req.window, dtype=np.float32)[np.newaxis, ...] - preds, inference_ms = _run_inference(input_array) - - return PredictResponse( - hip_left=float(preds[0]), - hip_right=float(preds[1]), - knee_left=float(preds[2]), - knee_right=float(preds[3]), - inference_ms=round(inference_ms, 3), - ) + return _run_inference(model_name, input_array) @app.post("/predict_msgpack") -async def predict_msgpack(request: Request) -> Response: +async def predict_msgpack(request: Request, model: str = DEFAULT_MODEL) -> Response: """ Msgpack endpoint for lower-overhead inference. - Request body: msgpack-encoded flat list of 187*28=5236 float32 values (row-major) - Response body: msgpack-encoded dict {hip_left, hip_right, knee_left, knee_right, inference_ms} + Query params: + model: Model name from MODEL_REGISTRY (default: 'full') + + Request body: + msgpack-encoded flat list of (window_size * feature_count) float32 values (row-major) + + Response body: + msgpack-encoded dict {"hip_left": , ..., "inference_ms": } """ - if _session is None or _input_name is None: - raise HTTPException(status_code=503, detail="Model not loaded") + model_name = _resolve_model(model) + cfg = MODEL_REGISTRY[model_name] + + if _targets_mean is None: + raise HTTPException(status_code=503, detail="Normalization stats not loaded") raw = await request.body() flat = msgpack.unpackb(raw, raw=False) - input_array = np.array(flat, dtype=np.float32).reshape(1, WINDOW_SIZE, FEATURE_COUNT) - preds, inference_ms = _run_inference(input_array) - - result = { - "hip_left": float(preds[0]), - "hip_right": float(preds[1]), - "knee_left": float(preds[2]), - "knee_right": float(preds[3]), - "inference_ms": round(inference_ms, 3), - } + expected_len = cfg["window_size"] * cfg["feature_count"] + if len(flat) != expected_len: + raise HTTPException( + status_code=422, + detail=( + f"Expected {expected_len} values for model '{model_name}' " + f"({cfg['window_size']} × {cfg['feature_count']}), got {len(flat)}" + ), + ) + + input_array = np.array(flat, dtype=np.float32).reshape( + 1, cfg["window_size"], cfg["feature_count"] + ) + result = _run_inference(model_name, input_array) return Response(content=msgpack.packb(result), media_type="application/x-msgpack") diff --git a/scripts/train_single_joint_tcn.py b/scripts/train_single_joint_tcn.py new file mode 100644 index 0000000..2d2cc89 --- /dev/null +++ b/scripts/train_single_joint_tcn.py @@ -0,0 +1,431 @@ +""" +Training script for single joint TCN model on exoskeleton joint moment estimation. + +Usage: + # Basic training with default config + python scripts/train_single_joint_tcn.py + + # Override specific parameters + python scripts/train_single_joint_tcn.py training.num_epochs=50 training.batch_size=16 + + # Use different model variant + python scripts/train_single_joint_tcn.py model=tcn_small + + # Hyperparameter sweep + python scripts/train_single_joint_tcn.py -m model.architecture.num_channels=[15,15,15],[25,25,25,25,25] +""" + +from exoskeleton_ml.utils import ( + EarlyStopping, + compute_metrics, + get_device, + save_best_model, + save_checkpoint, +) +from exoskeleton_ml.models import create_model +from exoskeleton_ml.data import create_dataloaders +import sys +from pathlib import Path + +import hydra +import torch +import torch.nn as nn +from omegaconf import DictConfig, OmegaConf +from torch.utils.data import DataLoader +from tqdm import tqdm + +# Add src to path +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + + +def masked_loss( + outputs: torch.Tensor, + targets: torch.Tensor, + mask: torch.Tensor, + criterion: nn.Module, +) -> torch.Tensor: + """Compute loss only on non-padded and non-NaN positions. + + Args: + outputs: Model predictions (batch, seq_len, 1). + targets: Ground truth (batch, seq_len, 1). + mask: Validity mask (batch, seq_len). + criterion: Loss function. + + Returns: + Masked loss value. + """ + + # Expand mask to match output dimensions: (batch, seq_len, 4) + mask_expanded = mask.unsqueeze(-1).expand_as(targets) + + # Create combined mask: valid positions AND non-NaN values in both outputs and targets + nan_mask_targets = ~torch.isnan(targets) + nan_mask_outputs = ~torch.isnan(outputs) + valid_mask = mask_expanded & nan_mask_targets & nan_mask_outputs + + # Convert boolean mask to float + valid_mask_float = valid_mask.float() + + # Compute squared error + squared_error = (outputs - targets) ** 2 + + # Apply mask to the squared error + masked_squared_error = squared_error * valid_mask_float + + # Count valid elements + num_valid = valid_mask_float.sum() + + # Check if we have any valid data + if num_valid == 0: + # Return zero loss if no valid data + return torch.tensor(0.0, device=outputs.device, requires_grad=True) + + # Compute mean squared error over valid positions only + loss = masked_squared_error.sum() / num_valid + + return loss + + +def train_epoch( + model: nn.Module, + loader: DataLoader, + criterion: nn.Module, + optimizer: torch.optim.Optimizer, + device: torch.device, + epoch: int, + cfg: DictConfig, +) -> float: + """Train for one epoch. + + Args: + model: PyTorch model. + loader: Training data loader. + criterion: Loss function. + optimizer: Optimizer. + device: Device to run on. + epoch: Current epoch number. + + Returns: + Average training loss for the epoch. + """ + model.train() + total_loss = 0.0 + num_batches = len(loader) + + progress_bar = tqdm(loader, desc=f"Epoch {epoch} [Train]", leave=False) + + for batch_idx, batch in enumerate(progress_bar): + + # obtain indicies from yaml configuration file + feature_idx = cfg.model.data_mapping.joint_features + joint_idx = cfg.model.data_mapping.joint_index + + inputs = batch["inputs"][:, :, feature_idx].to(device) # (batch, seq_len, 7) + targets = batch["targets"][:, :, joint_idx:joint_idx+1].to(device) # (batch, seq_len, 1) + mask = batch["mask"].to(device) # (batch, seq_len) + + # Forward pass + optimizer.zero_grad() + outputs = model(inputs) # (batch, seq_len, 4) + + # Compute masked loss + loss = masked_loss(outputs, targets, mask, criterion) + + # Backward pass + loss.backward() + + # Gradient clipping + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + + optimizer.step() + + # Update progress + total_loss += loss.item() + progress_bar.set_postfix({"loss": f"{loss.item():.4f}"}) + + avg_loss = total_loss / num_batches + return avg_loss + + +@torch.no_grad() +def validate( + model: nn.Module, + loader: DataLoader, + criterion: nn.Module, + device: torch.device, + epoch: int, + cfg: DictConfig, +) -> tuple[float, dict]: + """Validate on validation set. + + Args: + model: PyTorch model. + loader: Validation data loader. + criterion: Loss function. + device: Device to run on. + epoch: Current epoch number. + + Returns: + Tuple of (average loss, metrics dictionary). + """ + model.eval() + total_loss = 0.0 + all_outputs = [] + all_targets = [] + all_masks = [] + + # obtain indices from yaml configuration file + feature_idx = cfg.model.data_mapping.joint_features + joint_idx = cfg.model.data_mapping.joint_index + + + progress_bar = tqdm(loader, desc=f"Epoch {epoch} [Val]", leave=False) + + for batch in progress_bar: + inputs = batch["inputs"][:, :, feature_idx].to(device) + targets = batch["targets"][:, :, joint_idx:joint_idx+1].to(device) + mask = batch["mask"].to(device) + + # Forward pass + outputs = model(inputs) + + # Compute loss + loss = masked_loss(outputs, targets, mask, criterion) + total_loss += loss.item() + + # Store for metrics computation - flatten valid positions only + # This avoids issues with different sequence lengths across batches + for i in range(outputs.shape[0]): # Iterate over batch + seq_len = mask[i].sum().item() # Get actual length for this sequence + all_outputs.append(outputs[i, :seq_len].cpu()) # Only valid timesteps + all_targets.append(targets[i, :seq_len].cpu()) + all_masks.append(mask[i, :seq_len].cpu()) + + progress_bar.set_postfix({"loss": f"{loss.item():.4f}"}) + + # Compute average loss + avg_loss = total_loss / len(loader) + + # Concatenate all sequences (now all are 1D after flattening) + all_outputs = torch.cat(all_outputs, dim=0) # (total_valid_timesteps, 1) + all_targets = torch.cat(all_targets, dim=0) # (total_valid_timesteps, 1) + all_masks = torch.cat(all_masks, dim=0) # (total_valid_timesteps,) + + # Reshape for metrics computation - add batch and sequence dimensions back + # but now all sequences are concatenated into one long sequence + all_outputs = all_outputs.unsqueeze(0) # (1, total_timesteps, 1) + all_targets = all_targets.unsqueeze(0) # (1, total_timesteps, 1) + all_masks = all_masks.unsqueeze(0) # (1, total_timesteps) + + # Derive joint name from config for per-joint metrics key + _joint_label_map = ["hip_l", "hip_r", "knee_l", "knee_r"] + joint_name = _joint_label_map[cfg.model.data_mapping.joint_index] + + # Compute metrics using the single joint name so the returned dict has + # the correct key (e.g. "rmse_hip_l") instead of four 4-joint keys. + metrics = compute_metrics(all_outputs, all_targets, all_masks, joint_names=[joint_name]) + + return avg_loss, metrics + + +@hydra.main(config_path="../configs", config_name="config", version_base="1.3") +def train(cfg: DictConfig) -> None: + """Main training function. + + Args: + cfg: Hydra configuration object. + """ + print("=" * 80) + print("TCN Training for Exoskeleton Joint Moment Estimation") + print("=" * 80) + print("\nConfiguration:") + print(OmegaConf.to_yaml(cfg)) + + # Setup + device = get_device() + print(f"\nDevice: {device}") + + output_dir = Path(cfg.training.get("output_dir", "outputs/default")) + output_dir.mkdir(parents=True, exist_ok=True) + print(f"Output directory: {output_dir}") + + # Save config + with open(output_dir / "config.yaml", "w") as f: + OmegaConf.save(cfg, f) + + # Create dataloaders + print("\n" + "=" * 80) + print("Loading Data") + print("=" * 80) + + + train_loader, val_loader, test_loader = create_dataloaders( + hf_repo=cfg.data.hf_repo, + cache_dir=cfg.data.cache_dir, + train_participants=cfg.data.splits.train, + val_participants=cfg.data.splits.val, + test_participants=cfg.data.splits.test, + batch_size=cfg.training.batch_size, + num_workers=cfg.data.get("num_workers", 4), + normalize=cfg.data.preprocessing.get("normalize", True), + device=device.type, + ) + + print(f"Train batches: {len(train_loader)}") + print(f"Val batches: {len(val_loader)}") + print(f"Test batches: {len(test_loader)}") + + # Create model + print("\n" + "=" * 80) + print("Creating Model") + print("=" * 80) + + model = create_model(cfg.model).to(device) + num_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"Model: {cfg.model.name}") + print(f"Parameters: {num_params:,}") + + if hasattr(model, "get_effective_history"): + eff_hist = model.get_effective_history() + print( + f"Effective history: {eff_hist} timesteps ({eff_hist/100:.2f}s at 100Hz)") + + # Loss function + criterion = nn.MSELoss(reduction="mean") + + # Optimizer + optimizer = torch.optim.Adam( + model.parameters(), + lr=cfg.training.learning_rate, + weight_decay=cfg.training.get("weight_decay", 0.0001), + ) + + # Learning rate scheduler + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, + mode="min", + factor=0.5, + patience=10, + min_lr=1e-6, + ) + + # Early stopping + early_stopping = EarlyStopping( + patience=cfg.training.get("early_stopping_patience", 20), + min_delta=0.0, + mode="min", + ) + + # Training loop + print("\n" + "=" * 80) + print("Training") + print("=" * 80) + + # Derive joint label for display (mirrors logic in validate()) + _joint_label_map = ["hip_l", "hip_r", "knee_l", "knee_r"] + joint_label = _joint_label_map[cfg.model.data_mapping.joint_index] + + best_val_loss = float("inf") + train_losses = [] + val_losses = [] + + for epoch in range(1, cfg.training.num_epochs + 1): + print(f"\nEpoch {epoch}/{cfg.training.num_epochs}") + print("-" * 80) + + # Train + train_loss = train_epoch( + model, train_loader, criterion, optimizer, device, epoch, cfg) + train_losses.append(train_loss) + + # Validate + val_loss, val_metrics = validate( + model, val_loader, criterion, device, epoch, cfg) + val_losses.append(val_loss) + + # Print metrics + print(f"Train Loss: {train_loss:.4f}") + print(f"Val Loss: {val_loss:.4f}") + print(f"Val RMSE: {val_metrics['rmse_overall']:.4f} Nm/kg") + print(f"Val R²: {val_metrics['r2_overall']:.4f}") + print(f"Val MAE: {val_metrics['mae_overall']:.4f} Nm/kg") + print(f" {joint_label}: {val_metrics[f'rmse_{joint_label}']:.4f} Nm/kg") + + # Learning rate + current_lr = optimizer.param_groups[0]["lr"] + print(f"Learning Rate: {current_lr:.2e}") + + # Scheduler step + scheduler.step(val_loss) + + # Save best model + if val_loss < best_val_loss: + best_val_loss = val_loss + save_best_model( + model, + optimizer, + epoch, + val_loss, + output_dir, + metrics=val_metrics, + ) + print(f"✅ New best model saved! (Val Loss: {val_loss:.4f})") + + # Save periodic checkpoint + if epoch % cfg.training.get("save_frequency", 10) == 0: + save_checkpoint( + model, + optimizer, + epoch, + val_loss, + output_dir / f"checkpoint_epoch_{epoch}.pt", + scheduler=scheduler, + metrics=val_metrics, + ) + + # Early stopping check + early_stopping(val_loss) + if early_stopping.should_stop: + print(f"\n⚠️ Early stopping triggered at epoch {epoch}") + break + + # Final evaluation on test set + print("\n" + "=" * 80) + print("Final Evaluation on Test Set") + print("=" * 80) + + # Load best model + from exoskeleton_ml.utils import load_checkpoint + + load_checkpoint( + output_dir / "best_model.pt", + model, + device=str(device), + ) + + # Evaluate + test_loss, test_metrics = validate( + model, test_loader, criterion, device, 0, cfg) + + print(f"\nTest Loss: {test_loss:.4f}") + print(f"Test RMSE: {test_metrics['rmse_overall']:.4f} Nm/kg") + print(f"Test R²: {test_metrics['r2_overall']:.4f}") + print(f"Test MAE: {test_metrics['mae_overall']:.4f} Nm/kg") + print(f" {joint_label}: {test_metrics[f'rmse_{joint_label}']:.4f} Nm/kg") + + # Save final results + results = { + "best_val_loss": best_val_loss, + "test_loss": test_loss, + "test_metrics": test_metrics, + "train_losses": train_losses, + "val_losses": val_losses, + } + + torch.save(results, output_dir / "training_results.pt") + print(f"\n✅ Training complete! Results saved to {output_dir}") + + +if __name__ == "__main__": + train() diff --git a/src/exoskeleton_ml/utils/metrics.py b/src/exoskeleton_ml/utils/metrics.py index 7f9fbb6..dc20693 100644 --- a/src/exoskeleton_ml/utils/metrics.py +++ b/src/exoskeleton_ml/utils/metrics.py @@ -133,23 +133,29 @@ def compute_metrics( predictions: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor, + joint_names: list[str] | None = None, ) -> dict[str, float]: """Compute regression metrics for joint moment estimation. Args: - predictions: Model predictions of shape (batch, seq_len, 4). - targets: Ground truth targets of shape (batch, seq_len, 4). + predictions: Model predictions of shape (batch, seq_len, num_joints). + targets: Ground truth targets of shape (batch, seq_len, num_joints). mask: Boolean mask of shape (batch, seq_len) indicating valid positions. + joint_names: Names for each output joint. Defaults to + ["hip_l", "hip_r", "knee_l", "knee_r"] for the 4-joint model. Returns: Dictionary containing: - rmse_overall: Overall RMSE across all joints - - rmse_hip_l, rmse_hip_r, rmse_knee_l, rmse_knee_r: Per-joint RMSE + - rmse_: Per-joint RMSE for each joint name - mae_overall: Overall Mean Absolute Error - r2_overall: Overall R² score - nrmse_overall: Normalized RMSE (by target range) """ - # Expand mask to match predictions shape: (batch, seq_len, 1) + if joint_names is None: + joint_names = ["hip_l", "hip_r", "knee_l", "knee_r"] + + # Expand mask to match predictions shape: (batch, seq_len, num_joints) mask = mask.unsqueeze(-1) # Create combined mask: valid positions AND non-NaN values @@ -164,10 +170,7 @@ def compute_metrics( if num_valid == 0: return { "rmse_overall": 0.0, - "rmse_hip_l": 0.0, - "rmse_hip_r": 0.0, - "rmse_knee_l": 0.0, - "rmse_knee_r": 0.0, + **{f"rmse_{j}": 0.0 for j in joint_names}, "mae_overall": 0.0, "r2_overall": 0.0, "nrmse_overall": 0.0, @@ -178,7 +181,6 @@ def compute_metrics( rmse = torch.sqrt(mse).item() # Per-joint RMSE - joint_names = ["hip_l", "hip_r", "knee_l", "knee_r"] per_joint_rmse = {} for i, joint in enumerate(joint_names): # Get valid mask for this specific joint