diff --git a/docs/ml_control_interface.md b/docs/ml_control_interface.md new file mode 100644 index 0000000..24a430f --- /dev/null +++ b/docs/ml_control_interface.md @@ -0,0 +1,254 @@ +# ML–Control Interface Specification + +> **Version**: 1.0.0 +> **Status**: Draft — Pending Embedded Team Review +> **Last Updated**: 2026-01-21 + +--- + +## Overview + +This document defines the **interface contract** between the Machine Learning (Movement Prediction) module and the Control/Embedded system for the McMaster Exoskeleton project. + +**Purpose**: Enable independent development by both teams while ensuring integration compatibility. + +--- + +## Input Specification + +The ML model receives a sequence of sensor measurements collected by the embedded system. + +### Tensor Shape + +``` +(Batch_Size, Sequence_Length, 28) +``` + +| Dimension | Symbol | Description | +|-----------|--------|-------------| +| Batch Size | `B` | Number of parallel inference streams. **Real-time: 1** | +| Sequence Length | `L` | Variable length history (minimum ~100 samples recommended) | +| Input Features | `C_in` | **28** features (24 IMU + 4 encoder angles) | + +### Data Type + +- **dtype**: `float32` + +### Feature Ordering + +The 28 input features **must** be provided in the following order: + +#### IMU Features (Indices 0–23) + +| Index | Feature Name | Unit | Description | +|-------|--------------|------|-------------| +| 0 | `thigh_imu_l_accel_x` | m/s² | Left thigh IMU — Acceleration X | +| 1 | `thigh_imu_l_accel_y` | m/s² | Left thigh IMU — Acceleration Y | +| 2 | `thigh_imu_l_accel_z` | m/s² | Left thigh IMU — Acceleration Z | +| 3 | `thigh_imu_l_gyro_x` | rad/s | Left thigh IMU — Gyroscope X | +| 4 | `thigh_imu_l_gyro_y` | rad/s | Left thigh IMU — Gyroscope Y | +| 5 | `thigh_imu_l_gyro_z` | rad/s | Left thigh IMU — Gyroscope Z | +| 6 | `shank_imu_l_accel_x` | m/s² | Left shank IMU — Acceleration X | +| 7 | `shank_imu_l_accel_y` | m/s² | Left shank IMU — Acceleration Y | +| 8 | `shank_imu_l_accel_z` | m/s² | Left shank IMU — Acceleration Z | +| 9 | `shank_imu_l_gyro_x` | rad/s | Left shank IMU — Gyroscope X | +| 10 | `shank_imu_l_gyro_y` | rad/s | Left shank IMU — Gyroscope Y | +| 11 | `shank_imu_l_gyro_z` | rad/s | Left shank IMU — Gyroscope Z | +| 12 | `thigh_imu_r_accel_x` | m/s² | Right thigh IMU — Acceleration X | +| 13 | `thigh_imu_r_accel_y` | m/s² | Right thigh IMU — Acceleration Y | +| 14 | `thigh_imu_r_accel_z` | m/s² | Right thigh IMU — Acceleration Z | +| 15 | `thigh_imu_r_gyro_x` | rad/s | Right thigh IMU — Gyroscope X | +| 16 | `thigh_imu_r_gyro_y` | rad/s | Right thigh IMU — Gyroscope Y | +| 17 | `thigh_imu_r_gyro_z` | rad/s | Right thigh IMU — Gyroscope Z | +| 18 | `shank_imu_r_accel_x` | m/s² | Right shank IMU — Acceleration X | +| 19 | `shank_imu_r_accel_y` | m/s² | Right shank IMU — Acceleration Y | +| 20 | `shank_imu_r_accel_z` | m/s² | Right shank IMU — Acceleration Z | +| 21 | `shank_imu_r_gyro_x` | rad/s | Right shank IMU — Gyroscope X | +| 22 | `shank_imu_r_gyro_y` | rad/s | Right shank IMU — Gyroscope Y | +| 23 | `shank_imu_r_gyro_z` | rad/s | Right shank IMU — Gyroscope Z | + +#### Encoder Angle Features (Indices 24–27) + +| Index | Feature Name | Unit | Description | +|-------|--------------|------|-------------| +| 24 | `hip_flexion_l` | rad | Left hip encoder angle | +| 25 | `knee_angle_l` | rad | Left knee encoder angle | +| 26 | `hip_flexion_r` | rad | Right hip encoder angle | +| 27 | `knee_angle_r` | rad | Right knee encoder angle | + +--- + +## Output Specification + +The ML model outputs predicted biological joint moments. + +### Tensor Shape + +``` +(Batch_Size, Sequence_Length, 4) +``` + +For real-time control, use only the **last timestep**: `output[:, -1, :]` → shape `(1, 4)` + +### Data Type + +- **dtype**: `float32` + +### Feature Ordering + +| Index | Feature Name | Unit | Description | +|-------|--------------|------|-------------| +| 0 | `hip_moment_l` | Nm/kg | Left hip biological joint moment | +| 1 | `knee_moment_l` | Nm/kg | Left knee biological joint moment | +| 2 | `hip_moment_r` | Nm/kg | Right hip biological joint moment | +| 3 | `knee_moment_r` | Nm/kg | Right knee biological joint moment | + +> [!IMPORTANT] +> Output moments are **normalized by body mass** (Nm/kg). To obtain absolute torque (Nm), the embedded controller must multiply by the participant's mass in kg: +> ``` +> torque_nm = moment_nm_per_kg * participant_mass_kg +> ``` + +--- + +## Timing & Real-Time Constraints + +| Parameter | Value | Notes | +|-----------|-------|-------| +| **Sampling Rate** | 100 Hz | Sensor data collected every 10 ms | +| **Inference Latency** | < 5 ms | Target; allows margin for control loop overhead | +| **Model Causality** | Strictly Causal | Prediction at time *t* uses only inputs from *t* and earlier | +| **Minimum History** | ~100 samples | ~1 second of context for TCN receptive field | + +### Control Loop Timing Diagram + +``` +Time (ms): 0 10 20 30 40 + │ │ │ │ │ +Sensor: ├───┬───┼───┬───┼───┬───┼───┬───┤ + │ S₀│ │ S₁│ │ S₂│ │ S₃│ │ + └───┘ └───┘ └───┘ └───┘ + ↓ + ┌─────────────┐ + │ Inference │ < 5ms + └──────┬──────┘ + ↓ + ┌─────────────┐ + │ Actuation │ + └─────────────┘ +``` + +--- + +## Variable Sequence Length + +The TCN (Temporal Convolutional Network) architecture is a **Fully Convolutional Network**, which natively supports variable-length input sequences. + +### Training +- Sequences are padded to batch-max length +- A boolean mask indicates valid (non-padded) timesteps +- Loss is computed only on valid timesteps + +### Inference (Real-Time) +- Maintain a rolling buffer of recent samples +- No fixed window size required +- Recommended minimum buffer: 100 samples (~1 second) + +--- + +## Example Inference Call + +See [example_inference.py](./example_inference.py) for a complete working example. + +### Minimal Code Snippet + +```python +import torch + +# Load trained model +model = torch.load("models/tcn_moment_predictor.pt") +model.eval() + +# Prepare input: (1, sequence_length, 28) +sensor_buffer = torch.tensor(sensor_data, dtype=torch.float32).unsqueeze(0) + +# Run inference +with torch.no_grad(): + output = model(sensor_buffer) # (1, L, 4) + current_moment = output[:, -1, :] # (1, 4) — last timestep only + +# Denormalize to absolute torque +torque_nm = current_moment * participant_mass_kg +``` + +--- + +## Normalization + +### Input Normalization + + The model expects **normalized inputs**. Normalization statistics are computed from the training set and stored in: + + ``` + data/processed/phase1/normalization_stats.json + ``` + + Formula: + ``` + x_normalized = (x_raw - mean) / (std + 1e-8) + ``` + +> [!NOTE] +> The embedded system should either: +> 1. Apply normalization before sending data to the ML module, OR +> 2. Send raw data and let the ML module handle normalization + +### Output Denormalization + +If the model was trained with normalized targets, apply inverse transform: +``` +y_raw = y_normalized * std + mean +``` + +--- + +## Error Handling + +| Error Condition | Recommended Action | +|-----------------|-------------------| +| Input shape mismatch | Return error code; do not actuate | +| NaN/Inf in input | Clamp or skip inference cycle | +| Inference timeout (> 10ms) | Use previous prediction; log warning | +| Model load failure | Fall back to safe mode (zero torque) | + +--- + +## Versioning + +This interface follows semantic versioning: + +- **MAJOR**: Breaking changes to tensor shapes or feature ordering +- **MINOR**: Additions (new optional features) +- **PATCH**: Documentation or implementation fixes + +Current version: **1.0.0** + +--- + +## Open Questions + +> [!WARNING] +> The following items require confirmation from the controls/embedded team: + +1. **Sampling Rate**: Is 100 Hz confirmed, or could it change? +2. **Latency Budget**: Is < 5 ms acceptable? +3. **Denormalization Responsibility**: Should the ML module output absolute torque (Nm), or is Nm/kg acceptable? +4. **Joint Order**: Confirm `[hip_l, knee_l, hip_r, knee_r]` matches actuator wiring. +5. **IMU Frame Convention**: Confirm axis orientation (NED? ENU?). + +--- + +## References + +- [Data Infrastructure Plan](./data_infrastructure_plan.md) +- [ExoskeletonDataset Class](./src/exoskeleton_ml/data/datasets.py) diff --git a/pyproject.toml b/pyproject.toml index 85f6ea4..58068f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,9 @@ dependencies = [ "h5py>=3.9.0", "hydra-core>=1.3.0", "omegaconf>=2.3.0", + "onnx>=1.20.1", + "onnxscript>=0.5.7", + "onnxruntime>=1.23.2", "wandb>=0.16.0", "tensorboard>=2.15.0", "datasets>=2.14.0", diff --git a/scripts/sample_inference.py b/scripts/sample_inference.py index 41cc7bd..1692908 100644 --- a/scripts/sample_inference.py +++ b/scripts/sample_inference.py @@ -13,6 +13,8 @@ import numpy as np import torch +import onnx +import onnxruntime from exoskeleton_ml.models import create_model from exoskeleton_ml.utils import load_checkpoint @@ -30,7 +32,6 @@ def create_sample_input( def load_model(model_path: str, device: str = "cpu") -> torch.nn.Module: model_path = Path(model_path) checkpoint = torch.load(model_path, map_location=device) - config_path = model_path.parent / "config.yaml" if not config_path.exists(): raise ValueError(f"Config file not found at {config_path}") @@ -62,26 +63,120 @@ def run_inference( return predictions, inference_time_ms -def benchmark_inference( +def benchmark_pytorch( model: torch.nn.Module, input_data: torch.Tensor, num_runs: int = 100, + warmup_runs: int = 10, device: str = "cpu", ) -> dict: - print(f"\nRunning benchmark with {num_runs} iterations...") + """Benchmark PyTorch model inference.""" + print(f"\nBenchmarking PyTorch ({device})...") + input_data = input_data.to(device) + # Warmup + print(f" Warming up ({warmup_runs} runs)...") + for _ in range(warmup_runs): + with torch.no_grad(): + _ = model(input_data) + + # Benchmark + print(f" Benchmarking ({num_runs} runs)...") inference_times = [] + for _ in range(num_runs): + if device == "cuda": + torch.cuda.synchronize() + start = time.perf_counter() + with torch.no_grad(): + _ = model(input_data) + if device == "cuda": + torch.cuda.synchronize() + end = time.perf_counter() + inference_times.append((end - start) * 1000) - for i in range(num_runs): - _, inference_time = run_inference(model, input_data, device) - inference_times.append(inference_time) + inference_times = np.array(inference_times) + return { + "mean_ms": np.mean(inference_times), + "std_ms": np.std(inference_times), + "min_ms": np.min(inference_times), + "max_ms": np.max(inference_times), + "median_ms": np.median(inference_times), + "p95_ms": np.percentile(inference_times, 95), + "p99_ms": np.percentile(inference_times, 99), + } - if (i + 1) % 10 == 0: - print(f" Progress: {i + 1}/{num_runs}") + +def benchmark_onnx( + ort_session: onnxruntime.InferenceSession, + input_data: torch.Tensor, + num_runs: int = 100, + warmup_runs: int = 10, +) -> dict: + """Benchmark ONNX Runtime inference.""" + print(f"\nBenchmarking ONNX Runtime (CPU)...") + + input_name = ort_session.get_inputs()[0].name + onnx_input = {input_name: input_data.numpy()} + + # Warmup + print(f" Warming up ({warmup_runs} runs)...") + for _ in range(warmup_runs): + _ = ort_session.run(None, onnx_input) + + # Benchmark + print(f" Benchmarking ({num_runs} runs)...") + inference_times = [] + for _ in range(num_runs): + start = time.perf_counter() + _ = ort_session.run(None, onnx_input) + end = time.perf_counter() + inference_times.append((end - start) * 1000) inference_times = np.array(inference_times) + return { + "mean_ms": np.mean(inference_times), + "std_ms": np.std(inference_times), + "min_ms": np.min(inference_times), + "max_ms": np.max(inference_times), + "median_ms": np.median(inference_times), + "p95_ms": np.percentile(inference_times, 95), + "p99_ms": np.percentile(inference_times, 99), + } + + +def benchmark_torchscript( + scripted_model: torch.jit.ScriptModule, + input_data: torch.Tensor, + num_runs: int = 100, + warmup_runs: int = 10, + device: str = "cpu", +) -> dict: + """Benchmark TorchScript model inference.""" + print(f"\nBenchmarking TorchScript ({device})...") + input_data = input_data.to(device) + + # Warmup + print(f" Warming up ({warmup_runs} runs)...") + for _ in range(warmup_runs): + with torch.no_grad(): + _ = scripted_model(input_data) + + # Benchmark + print(f" Benchmarking ({num_runs} runs)...") + inference_times = [] + for _ in range(num_runs): + if device == "cuda": + torch.cuda.synchronize() + start = time.perf_counter() + with torch.no_grad(): + _ = scripted_model(input_data) + if device == "cuda": + torch.cuda.synchronize() + end = time.perf_counter() + inference_times.append((end - start) * 1000) - stats = { + inference_times = np.array(inference_times) + return { "mean_ms": np.mean(inference_times), "std_ms": np.std(inference_times), "min_ms": np.min(inference_times), @@ -91,7 +186,66 @@ def benchmark_inference( "p99_ms": np.percentile(inference_times, 99), } - return stats + +def print_comparison(pytorch_stats: dict, onnx_stats: dict, torchscript_stats: dict = None): + """Print comparison table between PyTorch, TorchScript, and ONNX.""" + print("\n" + "=" * 80) + print("TIMING COMPARISON: PyTorch vs TorchScript vs ONNX Runtime") + print("=" * 80) + + if torchscript_stats: + print(f"\n{'Metric':<10} {'PyTorch':<12} {'TorchScript':<14} {'ONNX':<12} {'TS Speedup':<12} {'ONNX Speedup':<12}") + print("-" * 75) + for metric, label in [("mean_ms", "Mean"), ("median_ms", "Median"), + ("min_ms", "Min"), ("max_ms", "Max"), + ("p95_ms", "P95"), ("p99_ms", "P99")]: + pt_val = pytorch_stats[metric] + ts_val = torchscript_stats[metric] + onnx_val = onnx_stats[metric] + ts_speedup = pt_val / ts_val if ts_val > 0 else 0 + onnx_speedup = pt_val / onnx_val if onnx_val > 0 else 0 + print(f"{label:<10} {pt_val:<12.3f} {ts_val:<14.3f} {onnx_val:<12.3f} {ts_speedup:.2f}x{'':<7} {onnx_speedup:.2f}x") + else: + print(f"\n{'Metric':<12} {'PyTorch (ms)':<15} {'ONNX (ms)':<15} {'Speedup':<10}") + print("-" * 55) + for metric, label in [("mean_ms", "Mean"), ("median_ms", "Median"), + ("min_ms", "Min"), ("max_ms", "Max"), + ("p95_ms", "P95"), ("p99_ms", "P99")]: + pt_val = pytorch_stats[metric] + onnx_val = onnx_stats[metric] + speedup = pt_val / onnx_val if onnx_val > 0 else 0 + print(f"{label:<12} {pt_val:<15.3f} {onnx_val:<15.3f} {speedup:.2f}x") + + pt_mean = pytorch_stats["mean_ms"] + onnx_mean = onnx_stats["mean_ms"] + ts_mean = torchscript_stats["mean_ms"] if torchscript_stats else None + + print("\n" + "-" * 75) + print("Summary:") + if ts_mean: + ts_speedup = pt_mean / ts_mean if ts_mean > 0 else 0 + if ts_speedup > 1: + print(f"TorchScript is {ts_speedup:.2f}x faster than PyTorch") + else: + print(f"PyTorch is {1/ts_speedup:.2f}x faster than TorchScript") + onnx_speedup = pt_mean / onnx_mean if onnx_mean > 0 else 0 + if onnx_speedup > 1: + print(f"ONNX is {onnx_speedup:.2f}x faster than PyTorch") + else: + print(f" ⚠️ PyTorch is {1/onnx_speedup:.2f}x faster than ONNX") + + # Find fastest + if ts_mean: + fastest = min([("PyTorch", pt_mean), ("TorchScript", ts_mean), ("ONNX", onnx_mean)], key=lambda x: x[1]) + print(f" ⭐ Fastest: {fastest[0]} ({fastest[1]:.2f}ms)") + + # Real-time check + target = 10.0 # 10ms for 100Hz + print(f"\nReal-time feasibility (target: <{target}ms for 100Hz):") + print(f" PyTorch: {'✅' if pt_mean < target else '❌'} {pt_mean:.2f}ms") + if ts_mean: + print(f" TorchScript: {'✅' if ts_mean < target else '❌'} {ts_mean:.2f}ms") + print(f" ONNX: {'✅' if onnx_mean < target else '❌'} {onnx_mean:.2f}ms") def main(): @@ -181,21 +335,56 @@ def main(): samples_per_second = 1000 / inference_time print(f"Throughput: {samples_per_second:.2f} sequences/second") + # Export to ONNX + print("\nExporting model to ONNX...") + onnx_path = "model.onnx" + onnx_program = torch.onnx.export(model, input_data, dynamo=True) + onnx_program.save(onnx_path) + onnx.checker.check_model(onnx_path) + print(f"✅ ONNX export successful: {onnx_path}") + + # Create ONNX session + session_options = onnxruntime.SessionOptions() + session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + ort_session = onnxruntime.InferenceSession( + onnx_path, sess_options=session_options, providers=["CPUExecutionProvider"] + ) + + # Verify outputs match + print("\nVerifying ONNX output matches PyTorch output...") + input_name = ort_session.get_inputs()[0].name + onnx_output = ort_session.run(None, {input_name: input_data.numpy()})[0] + try: + torch.testing.assert_close(predictions, torch.tensor(onnx_output), rtol=1e-3, atol=1e-4) + print("✅ Outputs match within tolerance") + except AssertionError as e: + print(f"❌ Outputs differ: {e}") + + # Export to TorchScript + print("\nExporting model to TorchScript...") + torchscript_path = "model_traced.pt" + scripted_model = torch.jit.trace(model, input_data) + scripted_model.save(torchscript_path) + print(f"✅ TorchScript export successful: {torchscript_path}") + + # Verify TorchScript output + print("\nVerifying TorchScript output matches PyTorch output...") + with torch.no_grad(): + ts_output = scripted_model(input_data) + try: + torch.testing.assert_close(predictions, ts_output, rtol=1e-5, atol=1e-6) + print("✅ TorchScript outputs match within tolerance") + except AssertionError as e: + print(f"❌ TorchScript outputs differ: {e}") + if args.benchmark: - stats = benchmark_inference( - model, input_data, num_runs=args.num_runs, device=args.device - ) - - print("\n" + "=" * 80) - print("Benchmark Results") - print("=" * 80) - print(f"Mean: {stats['mean_ms']:.2f} ms") - print(f"Median: {stats['median_ms']:.2f} ms") - print(f"Std: {stats['std_ms']:.2f} ms") - print(f"Min: {stats['min_ms']:.2f} ms") - print(f"Max: {stats['max_ms']:.2f} ms") - print(f"P95: {stats['p95_ms']:.2f} ms") - print(f"P99: {stats['p99_ms']:.2f} ms") + # Benchmark all three + pytorch_stats = benchmark_pytorch(model, input_data, num_runs=args.num_runs, device=args.device) + torchscript_stats = benchmark_torchscript(scripted_model, input_data, num_runs=args.num_runs, device=args.device) + onnx_stats = benchmark_onnx(ort_session, input_data, num_runs=args.num_runs) + print_comparison(pytorch_stats, onnx_stats, torchscript_stats) + + print("\n" + "=" * 80) print("Summary")