From 8f25361435e7869be8cc0b43f8de0dd0120a153c Mon Sep 17 00:00:00 2001 From: Evan Montgomery-Recht Date: Fri, 3 Jul 2026 18:20:14 -0400 Subject: [PATCH 1/5] fork: port tools/dgx to combined driving_supercombo.onnx MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upstream deleted the split driving_vision.onnx/driving_policy.onnx in the July 2026 restructure; the merged graph has 6 inputs (img, big_img, features_buffer (1,24,512), desire_pulse, traffic_convention, action_t) and one flat (1,2576) output, with metadata embedded in the ONNX instead of a sidecar pkl. - teacher.py: single supercombo TensorRT engine; real output parsing via embedded output_slices + modeld's Parser (replaces the placeholder that returned zeros); emits path_mean/std from the plan position MDN plus all parsed outputs and hidden_state features - train.py: student runs the full 6-input graph (batch-1 ONNX, per-sample); differentiable plan extraction feeds the Laplacian NLL — the old wiring passed a {raw: ...} dict the loss ignored, so total was always 0; dry-run now works (DummyTeacher instead of None); route-dataloader batch keys and 6->12 channel duplication handled - model_runner: dummy inputs from ONNX-embedded metadata (uint8 imgs), supercombo fallback spec - benchmarks: single supercombo pass replaces the two-stage pipeline; stale split-model tinygrad baselines dropped (numbers pending DGX re-validation) - tests: contract tests for the embedded metadata, output_slices coverage, Parser output shapes, and torch-vs-numpy plan extraction consistency - docs: ported to the combined flow, old benchmark numbers labeled historical Verified on CPU: metadata parse against the real ONNX, torch extraction matches Parser exactly, full training step produces finite non-zero loss with gradient flow. TensorRT/DoRA runs on DGX hardware still pending. Co-Authored-By: Claude Fable 5 --- .../tools/dgx/HARDWARE_TEST_CHECKLIST.md | 13 +- openpilot/tools/dgx/README.md | 36 ++-- openpilot/tools/dgx/TRAINING.md | 93 +++++++--- openpilot/tools/dgx/benchmark_inference.py | 67 ++----- openpilot/tools/dgx/benchmark_tensorrt.py | 17 +- openpilot/tools/dgx/model_runner.py | 52 +++--- .../tools/dgx/tests/test_model_inference.py | 100 +++++++++-- openpilot/tools/dgx/training/dataloader.py | 16 +- openpilot/tools/dgx/training/losses.py | 16 +- openpilot/tools/dgx/training/teacher.py | 170 ++++++++++-------- openpilot/tools/dgx/training/train.py | 115 ++++++++++-- 11 files changed, 450 insertions(+), 245 deletions(-) diff --git a/openpilot/tools/dgx/HARDWARE_TEST_CHECKLIST.md b/openpilot/tools/dgx/HARDWARE_TEST_CHECKLIST.md index ce2cc38f83e177..2563b62d456e09 100644 --- a/openpilot/tools/dgx/HARDWARE_TEST_CHECKLIST.md +++ b/openpilot/tools/dgx/HARDWARE_TEST_CHECKLIST.md @@ -62,10 +62,9 @@ print(f'EON: {HARDWARE.get_device_type()}') python openpilot/tools/dgx/benchmark_tensorrt.py --runs 50 --warmup 10 ``` Expected results (DGX Spark GB10): -- [ ] driving_policy: < 0.5ms (>2000 FPS) -- [ ] driving_vision: < 2ms (>500 FPS) +- [ ] driving_supercombo: < 2.5ms (target TBD on hardware; placeholder is the sum of the old split-model targets) - [ ] dmonitoring: < 1ms (>1000 FPS) -- [ ] Combined: < 3ms (>300 FPS) +- [ ] Combined: < 3.5ms (target TBD on hardware) - [ ] No CUDA errors or warnings ### 2.2 tinygrad CUDA Benchmark @@ -84,8 +83,7 @@ Device.DEFAULT = 'CUDA' # Attempt to load each model import onnx models = [ - 'openpilot/selfdrive/modeld/models/driving_policy.onnx', - 'openpilot/selfdrive/modeld/models/driving_vision.onnx', + 'openpilot/selfdrive/modeld/models/driving_supercombo.onnx', 'openpilot/selfdrive/modeld/models/dmonitoring_model.onnx', ] for m in models: @@ -98,6 +96,11 @@ for m in models: ``` - [ ] All models load successfully - [ ] No ONNX parsing errors +- [ ] Embedded supercombo metadata (input_shapes, output_slices) is readable via + `openpilot.selfdrive.modeld.get_model_metadata.make_metadata_dict(model_path)` + +Note: `big_driving_supercombo.onnx` (~1.75GB) is only used with comma's USB AMD GPU +and is not part of this checklist. ## 3. GPU Monitoring diff --git a/openpilot/tools/dgx/README.md b/openpilot/tools/dgx/README.md index fb417d9b38622f..ba05782eebf8f9 100644 --- a/openpilot/tools/dgx/README.md +++ b/openpilot/tools/dgx/README.md @@ -162,7 +162,7 @@ ssh -i ~/.ssh/dgx_key user@dgx-spark-ip ### Slow tinygrad performance tinygrad CUDA is not optimized for Blackwell architecture. Use TensorRT: ```bash -# TensorRT: ~800 FPS +# TensorRT: ~800 FPS (historical split-model measurement) python openpilot/tools/dgx/benchmark_tensorrt.py # tinygrad: ~2 FPS (not optimized) @@ -214,29 +214,41 @@ Options: - `--warmup N`: Number of warmup runs (default: 5) - `--beam N`: BEAM optimization level (default: 0, disabled) -### Benchmark Results (2026-01) +### Benchmark Results (2026-01, historical) -Tested on DGX Spark (GB10, Blackwell, compute 12.1): +Tested on DGX Spark (GB10, Blackwell, compute 12.1). + +**Note**: these numbers were measured against the pre-July-2026 split models +(`driving_vision.onnx` / `driving_policy.onnx`), which have been replaced upstream +by the combined `driving_supercombo.onnx`. The rows are kept as historical +reference; supercombo numbers are pending re-validation on DGX hardware. #### TensorRT FP16 (Recommended) +Historical measurements (pre-July-2026 split models): + | Model | Inference | FPS | vs tinygrad | |-------|-----------|-----|-------------| -| driving_policy | 0.09ms | 11,355 | 659x faster | -| driving_vision | 0.85ms | 1,181 | 432x faster | +| driving_policy (historical) | 0.09ms | 11,355 | 659x faster | +| driving_vision (historical) | 0.85ms | 1,181 | 432x faster | | dmonitoring | 0.28ms | 3,584 | 1,175x faster | -| **Combined** | **1.21ms** | **824** | **620x faster** | +| **Combined (historical)** | **1.21ms** | **824** | **620x faster** | +| driving_supercombo | pending re-validation on DGX hardware | — | — | -**TensorRT is 41x faster than comma 3X** (1.2ms vs ~50ms). +**TensorRT was 41x faster than comma 3X** (1.2ms vs ~50ms) in the historical +split-model benchmark. #### tinygrad CUDA (Not Optimized for Blackwell) +Historical measurements (pre-July-2026 split models): + | Model | Inference | FPS | |-------|-----------|-----| -| driving_policy | 58ms | 17.2 | -| driving_vision | 366ms | 2.7 | +| driving_policy (historical) | 58ms | 17.2 | +| driving_vision (historical) | 366ms | 2.7 | | dmonitoring | 328ms | 3.0 | -| **Combined** | **514ms** | **1.9** | +| **Combined (historical)** | **514ms** | **1.9** | +| driving_supercombo | pending re-validation on DGX hardware | — | tinygrad's CUDA backend is not yet optimized for Blackwell architecture. Use TensorRT for production-level performance. @@ -249,7 +261,7 @@ The DGX Spark is ideal for fine-tuning openpilot models using DoRA (Weight-Decom - **Parameter Efficient**: Only ~2-5% of parameters are trained - **Preserves Base Model**: Original weights are frozen, preventing catastrophic forgetting -- **Fast Training**: With TensorRT teacher, pseudo-label generation runs at 800+ FPS +- **Fast Training**: With TensorRT teacher, pseudo-label generation ran at 800+ FPS (historical split-model measurement; supercombo pending re-validation) - **Easy Deployment**: DoRA weights can be merged back into the base model for inference ### Quick Start Training @@ -273,7 +285,7 @@ python openpilot/tools/dgx/training/train.py --data comma_car_segments --epochs ```bash python openpilot/tools/dgx/training/train.py \ --data /path/to/segments \ # Training data path - --model openpilot/selfdrive/modeld/models/driving_policy.onnx \ + --model openpilot/selfdrive/modeld/models/driving_supercombo.onnx \ --epochs 10 \ # Number of training epochs --batch-size 32 \ # Batch size --lr 1e-4 \ # Learning rate diff --git a/openpilot/tools/dgx/TRAINING.md b/openpilot/tools/dgx/TRAINING.md index 141004d1226b22..53eb09ffab5d64 100644 --- a/openpilot/tools/dgx/TRAINING.md +++ b/openpilot/tools/dgx/TRAINING.md @@ -51,26 +51,47 @@ route_id/ **Extracted data per frame**: - Camera frame (YUV420 → RGB → normalized) -- Desire vector (8-dim: lane changes, turns) +- Desire vector (8-dim per frame: lane changes, turns; the model consumes a 25-frame `desire_pulse` history, shape (25, 8)) - Traffic convention (2-dim: left/right hand drive) - GPS coordinates (for filtering/validation) - Vehicle state (speed, steering angle, etc.) ### 2. Teacher Model (TensorRT) -Uses pre-trained comma models to generate pseudo-labels: +Uses the pre-trained comma supercombo model to generate pseudo-labels. Since the +July 2026 upstream restructure there is no separate vision stage: the combined +`driving_supercombo.onnx` does vision + policy in a single forward pass. ```python -# Teacher generates targets at 800+ FPS -teacher_vision = TensorRTEngine("driving_vision.onnx") -teacher_policy = TensorRTEngine("driving_policy.onnx") - -# For each batch of frames -features = teacher_vision(frames) # Visual features -targets = teacher_policy(features, desire) # Path predictions +# Teacher generates targets (800+ FPS measured historically on the split models) +teacher = TensorRTEngine("driving_supercombo.onnx") + +# Model metadata (input_shapes, output_shapes, output_slices, model_checkpoint) +# is embedded in the ONNX metadata_props -- there is no sidecar *_metadata.pkl +from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict +metadata = make_metadata_dict("driving_supercombo.onnx") +output_slices = metadata["output_slices"] + +# For each frame: one forward pass with the 6 supercombo inputs +outputs = teacher( + img=img, # (1, 12, 128, 256) uint8, road camera + big_img=big_img, # (1, 12, 128, 256) uint8, wide camera + features_buffer=features_buffer, # (1, 24, 512) float, rolling recurrent state + desire_pulse=desire_pulse, # (1, 25, 8) float + traffic_convention=traffic, # (1, 2) float + action_t=action_t, # (1, 2) float +) # -> single flat output vector (1, 2576) + +# Roll the recurrent state: this frame's hidden_state (512) feeds the next frame +hidden = outputs[:, output_slices["hidden_state"]] +features_buffer = np.concatenate([features_buffer[:, 1:], hidden[:, None]], axis=1) ``` -**Outputs (pseudo-labels)**: +**Outputs (pseudo-labels)**: the flat `outputs` vector is decomposed with +`output_slices` (meta, desire_pred, pose, wide_from_device_euler, road_transform, +lane_lines, lane_lines_prob, road_edges, lead, lead_prob, hidden_state, plan, +desire_state, pad) and parsed with `Parser().parse_outputs()` from +`openpilot.selfdrive.modeld.parse_model_outputs`: - Path predictions (33 future timestamps × 2 coords) - Lane line positions - Road edge positions @@ -103,7 +124,7 @@ class DoRALayer(nn.Module): ``` **Target layers for adaptation**: -- Final dense layers in policy network +- Final dense layers in the supercombo policy head - GRU hidden state projections - Output heads (paths, lanes, edges) @@ -140,27 +161,39 @@ def train_epoch(student, teacher, dataloader, optimizer, device): student.train() for batch in dataloader: - frames = batch['frames'].to(device) # (B, T, C, H, W) - desire = batch['desire'].to(device) # (B, T, 8) - traffic = batch['traffic'].to(device) # (B, 2) - - # Teacher generates pseudo-labels (no grad) + frames = batch['frames'].to(device) # (B, T, C, H, W) road camera + big_frames = batch['big_frames'].to(device) # (B, T, C, H, W) wide camera + desire_pulse = batch['desire_pulse'].to(device) # (B, T, 25, 8) + traffic = batch['traffic'].to(device) # (B, 2) + action_t = batch['action_t'].to(device) # (B, 2) + + # Teacher generates pseudo-labels (no grad): one supercombo pass per + # frame, rolling the previous frames' hidden_state through features_buffer with torch.no_grad(): - features = teacher.vision(frames) - targets = teacher.policy(features, desire, traffic) - - # Student forward pass - hidden = None + feat_buf = torch.zeros(frames.shape[0], 24, 512, device=device) + targets = [] + for t in range(frames.shape[1]): + out = teacher(frames[:, t], big_frames[:, t], feat_buf, + desire_pulse[:, t], traffic, action_t) + targets.append(out) + hidden = out[:, output_slices["hidden_state"]] # (B, 512) + feat_buf = torch.cat([feat_buf[:, 1:], hidden.unsqueeze(1)], dim=1) + + # Student forward pass (same rolling recurrent state) + student_buf = torch.zeros(frames.shape[0], 24, 512, device=device) total_loss = 0 for t in range(frames.shape[1]): pred, hidden = student( frames[:, t], - desire[:, t], + big_frames[:, t], + student_buf, + desire_pulse[:, t], traffic, - hidden + action_t, ) - loss = compute_loss(pred, targets[:, t]) + student_buf = torch.cat([student_buf[:, 1:], hidden.unsqueeze(1)], dim=1) + loss = compute_loss(pred, targets[t]) total_loss += loss # Backward pass (only DoRA params have gradients) @@ -188,10 +221,12 @@ def merge_dora_weights(model): def export_to_onnx(model, path): """Export merged model to ONNX for comma device.""" dummy_inputs = { - 'img': torch.randn(1, 12, 128, 256), - 'big_img': torch.randn(1, 12, 128, 256), - 'desire': torch.randn(1, 8), + 'img': torch.randint(0, 256, (1, 12, 128, 256), dtype=torch.uint8), + 'big_img': torch.randint(0, 256, (1, 12, 128, 256), dtype=torch.uint8), + 'features_buffer': torch.randn(1, 24, 512), + 'desire_pulse': torch.randn(1, 25, 8), 'traffic_convention': torch.randn(1, 2), + 'action_t': torch.randn(1, 2), } torch.onnx.export(model, dummy_inputs, path, opset_version=14) ``` @@ -225,7 +260,7 @@ openpilot/tools/dgx/ DGX Spark advantages: - 128GB unified memory = large batch sizes -- TensorRT teacher at 800+ FPS = fast pseudo-label generation +- TensorRT teacher at 800+ FPS (historical split-model measurement) = fast pseudo-label generation - Single device = no distributed training complexity ## Usage @@ -245,7 +280,7 @@ python openpilot/tools/dgx/training/train.py \ # 3. Export to ONNX python openpilot/tools/dgx/training/export.py \ --checkpoint best_model.pt \ - --output custom_driving_policy.onnx + --output custom_driving_supercombo.onnx ``` ## Next Steps diff --git a/openpilot/tools/dgx/benchmark_inference.py b/openpilot/tools/dgx/benchmark_inference.py index ef1c5c6b19f595..bfcc5de512b871 100755 --- a/openpilot/tools/dgx/benchmark_inference.py +++ b/openpilot/tools/dgx/benchmark_inference.py @@ -54,66 +54,27 @@ def main(): models_dir = "openpilot/selfdrive/modeld/models" - # driving_policy - print("\n[driving_policy.onnx]") - policy = OnnxRunner(f"{models_dir}/driving_policy.onnx") - policy_inputs = { - "desire_pulse": Tensor(np.random.randn(1, 25, 8).astype(np.float16)), - "traffic_convention": Tensor(np.random.randn(1, 2).astype(np.float16)), - "features_buffer": Tensor(np.random.randn(1, 25, 512).astype(np.float16)), - } - stats = benchmark_model(policy, policy_inputs, args.warmup, args.runs) - print(f" {stats['mean_ms']:.2f}ms +/- {stats['std_ms']:.2f}ms ({stats['fps']:.1f} FPS)") - - # driving_vision - print("\n[driving_vision.onnx]") - vision = OnnxRunner(f"{models_dir}/driving_vision.onnx") - vision_inputs = { - "img": Tensor(np.random.randint(0, 255, (1, 12, 128, 256), dtype=np.uint8)), - "big_img": Tensor(np.random.randint(0, 255, (1, 12, 128, 256), dtype=np.uint8)), - } - stats = benchmark_model(vision, vision_inputs, args.warmup, args.runs) + def random_inputs(runner) -> dict: + inputs = {} + for k, v in runner.get_empty_input_data().items(): + if "float" in str(v.dtype): + inputs[k] = Tensor(np.random.randn(*v.shape).astype(np.float32)) + else: + inputs[k] = Tensor(np.random.randint(0, 255, v.shape, dtype=np.uint8)) + return inputs + + # driving_supercombo: merged vision+policy graph, one pass does everything + print("\n[driving_supercombo.onnx]") + supercombo = OnnxRunner(f"{models_dir}/driving_supercombo.onnx") + stats = benchmark_model(supercombo, random_inputs(supercombo), args.warmup, args.runs) print(f" {stats['mean_ms']:.2f}ms +/- {stats['std_ms']:.2f}ms ({stats['fps']:.1f} FPS)") # dmonitoring_model print("\n[dmonitoring_model.onnx]") dmon = OnnxRunner(f"{models_dir}/dmonitoring_model.onnx") - dmon_empty = dmon.get_empty_input_data() - dmon_inputs = {} - for k, v in dmon_empty.items(): - if "float" in str(v.dtype): - dmon_inputs[k] = Tensor(np.random.randn(*v.shape).astype(np.float32)) - else: - dmon_inputs[k] = Tensor(np.random.randint(0, 255, v.shape, dtype=np.uint8)) - stats = benchmark_model(dmon, dmon_inputs, args.warmup, args.runs) + stats = benchmark_model(dmon, random_inputs(dmon), args.warmup, args.runs) print(f" {stats['mean_ms']:.2f}ms +/- {stats['std_ms']:.2f}ms ({stats['fps']:.1f} FPS)") - # Combined pipeline - print("\n[Combined Vision + Policy Pipeline]") - times = [] - for _ in range(args.warmup): - vout = vision(vision_inputs) - for v in vout.values(): - v.realize() - pout = policy(policy_inputs) - for v in pout.values(): - v.realize() - - for _ in range(args.runs): - start = time.perf_counter() - vout = vision(vision_inputs) - for v in vout.values(): - v.realize() - pout = policy(policy_inputs) - for v in pout.values(): - v.realize() - times.append(time.perf_counter() - start) - - mean_ms = np.mean(times) * 1000 - std_ms = np.std(times) * 1000 - fps = 1 / np.mean(times) - print(f" {mean_ms:.2f}ms +/- {std_ms:.2f}ms ({fps:.1f} FPS)") - print("\n" + "=" * 60) print("Benchmark complete!") diff --git a/openpilot/tools/dgx/benchmark_tensorrt.py b/openpilot/tools/dgx/benchmark_tensorrt.py index 2dca971c6dc804..e69a0590abda38 100755 --- a/openpilot/tools/dgx/benchmark_tensorrt.py +++ b/openpilot/tools/dgx/benchmark_tensorrt.py @@ -109,8 +109,7 @@ def main(): models_dir = "openpilot/selfdrive/modeld/models" models = [ - ("driving_policy.onnx", f"{models_dir}/driving_policy.onnx"), - ("driving_vision.onnx", f"{models_dir}/driving_vision.onnx"), + ("driving_supercombo.onnx", f"{models_dir}/driving_supercombo.onnx"), ("dmonitoring_model.onnx", f"{models_dir}/dmonitoring_model.onnx"), ] @@ -143,21 +142,19 @@ def main(): total = sum(results.values()) print(f"[Combined Pipeline]: {total:.3f}ms ({1000 / total:.1f} FPS)") - # Comparison with tinygrad + # Comparison with tinygrad CUDA (measure via benchmark_inference.py; the old + # split-model numbers don't apply to driving_supercombo) print("\n[Comparison with tinygrad CUDA]") - tinygrad_times = { - "driving_policy.onnx": 58, - "driving_vision.onnx": 366, - "dmonitoring_model.onnx": 328, + tinygrad_times: dict[str, float] = { + "dmonitoring_model.onnx": 328, # historical (2026-01, Blackwell) } for name, ms in results.items(): tinygrad_ms = tinygrad_times.get(name, 0) if tinygrad_ms > 0: speedup = tinygrad_ms / ms print(f" {name}: {speedup:.0f}x faster") - - tinygrad_total = sum(tinygrad_times.values()) - print(f"\n Combined: {tinygrad_total / total:.0f}x faster than tinygrad CUDA") + else: + print(f" {name}: no tinygrad baseline yet — run benchmark_inference.py to measure") # Comparison with comma 3X comma3x_ms = 50 # ~20 FPS target diff --git a/openpilot/tools/dgx/model_runner.py b/openpilot/tools/dgx/model_runner.py index 8b4f5abdbffbaa..812b8461cdd27c 100644 --- a/openpilot/tools/dgx/model_runner.py +++ b/openpilot/tools/dgx/model_runner.py @@ -105,7 +105,7 @@ class ModelRunner: Usage: runner = ModelRunner( - model_path="openpilot/selfdrive/modeld/models/driving_policy.onnx", + model_path="openpilot/selfdrive/modeld/models/driving_supercombo.onnx", backend=Backend.TENSORRT, precision=Precision.FP16, ) @@ -333,28 +333,40 @@ def _run_onnxruntime(self, inputs: dict[str, np.ndarray]) -> np.ndarray: def _create_dummy_inputs(self) -> dict[str, np.ndarray]: """Create dummy inputs for warmup.""" - # Load metadata if available - metadata_path = self.model_path.with_name(self.model_path.stem + "_metadata.pkl") - if metadata_path.exists(): - with open(metadata_path, "rb") as f: - metadata = pickle.load(f) - self.input_shapes = metadata.get("input_shapes", {}) + # Driving models embed their metadata in the ONNX (there is no sidecar + # pkl for them anymore); dmonitoring still ships a *_metadata.pkl + if self.model_path.suffix == ".onnx": + try: + from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict + + self.input_shapes = make_metadata_dict(self.model_path)["input_shapes"] + except Exception: # no embedded metadata / unparseable — fall through + pass if not self.input_shapes: - # Default shapes for openpilot models - if "vision" in str(self.model_path): - self.input_shapes = { - "input_img": (1, 12, 128, 256), - "calib": (1, 3), - } - elif "policy" in str(self.model_path): - self.input_shapes = { - "input": (1, 512), - } - else: - self.input_shapes = {"input": (1, 256)} + metadata_path = self.model_path.with_name(self.model_path.stem + "_metadata.pkl") + if metadata_path.exists(): + with open(metadata_path, "rb") as f: + metadata = pickle.load(f) + self.input_shapes = metadata.get("input_shapes", {}) - return {name: np.random.randn(*shape).astype(np.float32) for name, shape in self.input_shapes.items()} + if not self.input_shapes: + # Fallback: supercombo input spec + self.input_shapes = { + "img": (1, 12, 128, 256), + "big_img": (1, 12, 128, 256), + "desire_pulse": (1, 25, 8), + "traffic_convention": (1, 2), + "features_buffer": (1, 24, 512), + "action_t": (1, 2), + } + + def dummy(name: str, shape: tuple[int, ...]) -> np.ndarray: + if "img" in name: # camera inputs are uint8 YUV + return np.random.randint(0, 255, shape, dtype=np.uint8) + return np.random.randn(*shape).astype(np.float32) + + return {name: dummy(name, shape) for name, shape in self.input_shapes.items()} def profile( self, diff --git a/openpilot/tools/dgx/tests/test_model_inference.py b/openpilot/tools/dgx/tests/test_model_inference.py index 2fb6a9515744f7..4b0e6dd21a803f 100644 --- a/openpilot/tools/dgx/tests/test_model_inference.py +++ b/openpilot/tools/dgx/tests/test_model_inference.py @@ -92,7 +92,7 @@ def test_runner_initialization(self): from openpilot.tools.dgx.model_runner import Backend, ModelRunner, Precision runner = ModelRunner( - model_path="openpilot/selfdrive/modeld/models/driving_policy.onnx", + model_path="openpilot/selfdrive/modeld/models/driving_supercombo.onnx", backend=Backend.CPU, # Use CPU for init test precision=Precision.FP32, ) @@ -191,20 +191,96 @@ def test_model_files_exist(self, model_dir: Path): assert len(onnx_files) > 0 or len(pkl_files) > 0, f"No model files found in {model_dir}" - @requires_gpu - def test_load_driving_policy_metadata(self, model_dir: Path): - """Test loading driving policy metadata.""" - metadata_path = model_dir / "driving_policy_metadata.pkl" - if not metadata_path.exists(): - pytest.skip("Metadata file not found (run model compilation first)") +class TestSupercomboContract: + """Contract tests against the combined driving_supercombo.onnx. + + The driving model's metadata (input_shapes, output_slices, ...) is embedded + in the ONNX metadata_props — there is no sidecar *_metadata.pkl anymore. + These run on CPU: metadata extraction only parses the protobuf header, and + the parse tests use random vectors, never the model itself. + """ + + MODEL_PATH = Path("openpilot/selfdrive/modeld/models/driving_supercombo.onnx") - import pickle + EXPECTED_INPUTS = {"img", "big_img", "desire_pulse", "traffic_convention", "features_buffer", "action_t"} + EXPECTED_SLICES = { + "plan", "lane_lines", "lane_lines_prob", "road_edges", "lead", "lead_prob", + "meta", "desire_state", "desire_pred", "pose", "wide_from_device_euler", + "road_transform", "hidden_state", + } - with open(metadata_path, "rb") as f: - metadata = pickle.load(f) + @pytest.fixture + def metadata(self) -> dict: + if not self.MODEL_PATH.exists(): + pytest.skip(f"{self.MODEL_PATH} not found") + from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict + + try: + return make_metadata_dict(self.MODEL_PATH) + except Exception as e: + pytest.skip(f"could not parse embedded metadata (git-lfs pointer not pulled?): {e}") - assert "input_shapes" in metadata - assert "output_shapes" in metadata or "output_slices" in metadata + @pytest.fixture + def output_len(self, metadata: dict) -> int: + return next(iter(metadata["output_shapes"].values()))[-1] + + def test_embedded_metadata_inputs(self, metadata: dict): + from openpilot.selfdrive.modeld.constants import ModelConstants + + input_shapes = metadata["input_shapes"] + assert set(input_shapes) == self.EXPECTED_INPUTS + assert input_shapes["img"] == input_shapes["big_img"] + assert input_shapes["img"][0] == 1 and input_shapes["img"][1] == 6 * ModelConstants.N_FRAMES + assert input_shapes["features_buffer"][-1] == ModelConstants.FEATURE_LEN + assert input_shapes["desire_pulse"][-1] == ModelConstants.DESIRE_LEN + assert input_shapes["traffic_convention"] == (1, ModelConstants.TRAFFIC_CONVENTION_LEN) + + def test_output_slices_cover_output(self, metadata: dict, output_len: int): + from openpilot.selfdrive.modeld.constants import ModelConstants + + slices = metadata["output_slices"] + assert self.EXPECTED_SLICES <= set(slices) + hidden = slices["hidden_state"] + assert hidden.stop - hidden.start == ModelConstants.FEATURE_LEN + for name, sl in slices.items(): + lo, hi, _ = sl.indices(output_len) + assert 0 <= lo <= hi <= output_len, f"slice {name}={sl} outside output of len {output_len}" + + def test_parse_supercombo_outputs(self, metadata: dict, output_len: int): + from openpilot.selfdrive.modeld.constants import ModelConstants + from openpilot.tools.dgx.training.teacher import parse_supercombo_outputs + + rng = np.random.default_rng(0) + raw = rng.standard_normal((2, output_len), dtype=np.float32) + parsed = parse_supercombo_outputs(raw, metadata["output_slices"]) + + assert parsed["plan"].shape == (2, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + assert parsed["plan_stds"].shape == parsed["plan"].shape + assert np.all(parsed["plan_stds"] > 0) + assert parsed["lane_lines"].shape == (2, ModelConstants.NUM_LANE_LINES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH) + assert parsed["road_edges"].shape == (2, ModelConstants.NUM_ROAD_EDGES, ModelConstants.IDX_N, ModelConstants.ROAD_EDGES_WIDTH) + assert parsed["desire_state"].shape == (2, ModelConstants.DESIRE_PRED_WIDTH) + np.testing.assert_allclose(parsed["desire_state"].sum(axis=-1), 1.0, atol=1e-5) # softmaxed + # hidden_state passes through unparsed, and parsing must not mutate raw + np.testing.assert_array_equal(parsed["hidden_state"], raw[:, metadata["output_slices"]["hidden_state"]]) + + def test_extract_path_distribution_matches_parser(self, metadata: dict, output_len: int): + """Torch (student) path extraction agrees with the numpy (teacher) parser.""" + torch = pytest.importorskip("torch") + from openpilot.selfdrive.modeld.constants import Plan + from openpilot.tools.dgx.training.teacher import parse_supercombo_outputs + from openpilot.tools.dgx.training.train import extract_path_distribution + + rng = np.random.default_rng(1) + raw = rng.standard_normal((2, output_len), dtype=np.float32) + + mean, std = extract_path_distribution(torch.from_numpy(raw.copy()), metadata["output_slices"]) + parsed = parse_supercombo_outputs(raw, metadata["output_slices"]) + + assert mean.shape == (2, 1, 33, 3) and std.shape == (2, 1, 33, 3) + assert bool((std > 0).all()) + np.testing.assert_allclose(mean.numpy()[:, 0], parsed["plan"][:, :, Plan.POSITION], rtol=1e-5) + np.testing.assert_allclose(std.numpy()[:, 0], parsed["plan_stds"][:, :, Plan.POSITION], rtol=1e-5) class TestMemoryUtils: diff --git a/openpilot/tools/dgx/training/dataloader.py b/openpilot/tools/dgx/training/dataloader.py index 1008ac2fd24f7a..1798546bfbdcf3 100644 --- a/openpilot/tools/dgx/training/dataloader.py +++ b/openpilot/tools/dgx/training/dataloader.py @@ -51,8 +51,10 @@ class TrainingSample: desire: np.ndarray # (8,) one-hot desire vector traffic_convention: np.ndarray # (2,) left/right hand drive - # Labels from recorded model outputs - model_outputs: np.ndarray | None # (1000,) raw model output for distillation + # Legacy placeholder assembled from parsed modelV2 logs; its layout matches + # neither the old split models nor driving_supercombo's output_slices. + # Distillation labels come from TeacherModel.generate_labels at train time. + model_outputs: np.ndarray | None # (1000,) placeholder, unused by train.py # Metadata timestamp: int @@ -170,10 +172,12 @@ def _load_sample(self, seg_id: str, frame_idx: int) -> TrainingSample: ) def _extract_model_output(self, model_data) -> np.ndarray: - """Extract raw model output from modelV2 message.""" - # The modelV2 message contains parsed predictions - # For distillation, we want the raw output vector - # This is a simplified version - full implementation needs output format details + """Extract a coarse feature vector from the parsed modelV2 message. + + Legacy placeholder: modelV2 logs carry parsed predictions, not the raw + output vector, so this layout matches no model's output_slices. Training + distills against TeacherModel.generate_labels instead. + """ outputs = [] # Path predictions diff --git a/openpilot/tools/dgx/training/losses.py b/openpilot/tools/dgx/training/losses.py index 1826be8686811b..1a119494130149 100644 --- a/openpilot/tools/dgx/training/losses.py +++ b/openpilot/tools/dgx/training/losses.py @@ -35,9 +35,11 @@ def forward( """Compute Laplacian NLL loss. Args: - pred_mean: (batch, num_hypotheses, horizon, 2) predicted positions - pred_std: (batch, num_hypotheses, horizon, 2) predicted uncertainties - target: (batch, horizon, 2) ground truth positions + pred_mean: (batch, num_hypotheses, horizon, coords) predicted positions. + The supercombo plan head is a single-hypothesis MDN, so + num_hypotheses is 1 there (winner-takes-all degenerates gracefully). + pred_std: (batch, num_hypotheses, horizon, coords) predicted uncertainties + target: (batch, horizon, coords) ground truth positions Returns: Scalar loss value @@ -122,9 +124,11 @@ def forward( Args: pred: Dictionary with keys: - - 'path_mean': (batch, num_hyp, horizon, 2) - - 'path_std': (batch, num_hyp, horizon, 2) - - 'path_prob': (batch, num_hyp) hypothesis probabilities + - 'path_mean': (batch, num_hyp, horizon, coords); num_hyp=1 for the + single-hypothesis supercombo plan head + - 'path_std': (batch, num_hyp, horizon, coords) + - 'path_prob': (batch, num_hyp) optional; skipped when either side + omits it (no hypothesis weights in a single-hypothesis head) - 'lane_lines': (batch, num_lanes, horizon, 2) optional - 'road_edges': (batch, 2, horizon, 2) optional target: Dictionary with same structure from teacher diff --git a/openpilot/tools/dgx/training/teacher.py b/openpilot/tools/dgx/training/teacher.py index 7d67bdd9a3638d..8f2069a7c91e39 100644 --- a/openpilot/tools/dgx/training/teacher.py +++ b/openpilot/tools/dgx/training/teacher.py @@ -10,6 +10,10 @@ import numpy as np +from openpilot.selfdrive.modeld.constants import Plan +from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict +from openpilot.selfdrive.modeld.parse_model_outputs import Parser + # Type hints for optional imports try: import tensorrt as trt # type: ignore[import-not-found] @@ -107,12 +111,36 @@ def __call__(self, **inputs) -> dict[str, np.ndarray]: return {name: buf.copy() for name, buf in self.outputs.items()} +def parse_supercombo_outputs(raw_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]: + """Decompose and parse the flat supercombo output vector. + + Mirrors modeld's slice_outputs + Parser().parse_outputs flow + (openpilot/selfdrive/modeld/modeld.py), batched. + + Args: + raw_outputs: (batch, N) flat combined model output + output_slices: name -> slice mapping from the ONNX-embedded metadata + + Returns: + Parsed outputs keyed by name (plan, plan_stds, lane_lines, road_edges, + lead, meta, desire_state, hidden_state, ...). Values keep the batch dim. + """ + # Parser mutates in place (softmax/exp on views), so slice copies + sliced = {k: raw_outputs[:, v].copy() for k, v in output_slices.items()} + return Parser().parse_outputs(sliced) + + class TeacherModel: - """Combined vision + policy teacher for pseudo-label generation.""" + """Combined supercombo teacher for pseudo-label generation. + + Runs the single driving_supercombo.onnx (merged vision+policy graph); + input/output specs come from the metadata embedded in the ONNX. + """ def __init__( self, models_dir: str = "openpilot/selfdrive/modeld/models", + model_name: str = "driving_supercombo.onnx", fp16: bool = True, verbose: bool = False, ): @@ -120,19 +148,17 @@ def __init__( self.fp16 = fp16 self.verbose = verbose - # Load models - print("Loading teacher models with TensorRT...") - self.vision = self._load_model("driving_vision.onnx") - self.policy = self._load_model("driving_policy.onnx") - print("Teacher models loaded!") + model_path = self.models_dir / model_name + if not model_path.exists(): + raise FileNotFoundError(f"Model not found: {model_path}") - def _load_model(self, name: str) -> TensorRTEngine: - path = self.models_dir / name - if not path.exists(): - raise FileNotFoundError(f"Model not found: {path}") + self.metadata = make_metadata_dict(model_path) + self.input_shapes: dict[str, tuple[int, ...]] = self.metadata["input_shapes"] + self.output_slices: dict[str, slice] = self.metadata["output_slices"] - print(f" Building {name}...") - return TensorRTEngine(str(path), fp16=self.fp16, verbose=self.verbose) + print(f"Loading teacher model with TensorRT...\n Building {model_name}...") + self.model = TensorRTEngine(str(model_path), fp16=self.fp16, verbose=self.verbose) + print("Teacher model loaded!") def generate_labels( self, @@ -140,56 +166,70 @@ def generate_labels( big_img: np.ndarray, desire: np.ndarray, traffic_convention: np.ndarray, + features_buffer: np.ndarray | None = None, + action_t: np.ndarray | None = None, ) -> dict[str, np.ndarray]: """Generate pseudo-labels for a batch of frames. + The engine is built for batch 1 (the ONNX has fixed shapes), so frames + are run one at a time. Without an explicit features_buffer the recurrent + feature history is zero (cold start) — fine for shuffled training frames, + but sequential streams get better labels by threading each frame's + returned features back in. + Args: - img: (batch, 12, 128, 256) uint8 camera frames + img: (batch, 12, 128, 256) uint8 camera frames (2 stacked YUV frames) big_img: (batch, 12, 128, 256) uint8 wide camera frames - desire: (batch, 8) float16 desire vector - traffic_convention: (batch, 2) float16 traffic convention + desire: (batch, 8) desire vector, placed in the last desire_pulse step + traffic_convention: (batch, 2) traffic convention + features_buffer: (batch, 24, 512) prior feature history, zeros if None + action_t: (batch, 2) previous action, zeros if None Returns: Dictionary with: - - features: (batch, 512) visual features - - path_mean: (batch, num_hyp, horizon, 2) path predictions - - path_std: (batch, num_hyp, horizon, 2) uncertainties - - path_prob: (batch, num_hyp) hypothesis probabilities + - features: (batch, 512) hidden_state features (recurrent feedback) + - raw_outputs: (batch, N) flat model output + - path_mean: (batch, 1, 33, 3) plan position mean (single hypothesis) + - path_std: (batch, 1, 33, 3) plan position std + - path_prob: (batch, 1) hypothesis probability (always 1) + - all parsed outputs (plan, plan_stds, lane_lines, road_edges, lead, ...) """ - # Vision model: extract features - vision_out = self.vision(img=img, big_img=big_img) - - # Get features (assumed to be in vision output) - features = vision_out.get("features", vision_out.get("output", None)) - if features is None: - # Take first output if key not found - features = list(vision_out.values())[0] - - # Policy model: generate predictions - # Build features buffer for policy (1, 25, 512) from single frame features batch_size = img.shape[0] - features_buffer = np.zeros((batch_size, 25, 512), dtype=np.float16) - features_buffer[:, -1, :] = features.reshape(batch_size, -1)[:, :512] - - # Desire pulse (1, 25, 8) from single desire - desire_pulse = np.zeros((batch_size, 25, 8), dtype=np.float16) - desire_pulse[:, -1, :] = desire - - policy_out = self.policy( - desire_pulse=desire_pulse, - traffic_convention=traffic_convention, - features_buffer=features_buffer, - ) - - # Parse policy outputs - # Output shape is (1, 1000) - need to parse into structured predictions - outputs = policy_out.get("outputs", list(policy_out.values())[0]) - + pulse_shape = self.input_shapes["desire_pulse"][1:] # (25, 8) + feat_shape = self.input_shapes["features_buffer"][1:] # (24, 512) + + if features_buffer is None: + features_buffer = np.zeros((batch_size, *feat_shape), dtype=np.float32) + if action_t is None: + action_t = np.zeros((batch_size, self.input_shapes["action_t"][-1]), dtype=np.float32) + + raw = [] + for i in range(batch_size): + desire_pulse = np.zeros((1, *pulse_shape), dtype=np.float32) + desire_pulse[0, -1, :] = desire[i] + + out = self.model( + img=img[i : i + 1], + big_img=big_img[i : i + 1], + desire_pulse=desire_pulse, + traffic_convention=traffic_convention[i : i + 1], + features_buffer=features_buffer[i : i + 1], + action_t=action_t[i : i + 1], + ) + raw.append(out.get("outputs", next(iter(out.values()))).reshape(1, -1)) + + raw_outputs = np.concatenate(raw, axis=0).astype(np.float32) + parsed = parse_supercombo_outputs(raw_outputs, self.output_slices) + + plan = parsed["plan"] # (batch, 33, 15) + plan_stds = parsed["plan_stds"] return { - "features": features, - "raw_outputs": outputs, - # TODO: Parse outputs into path_mean, path_std, path_prob - # This requires understanding the exact output format + **parsed, + "features": raw_outputs[:, self.output_slices["hidden_state"]], + "raw_outputs": raw_outputs, + "path_mean": plan[:, np.newaxis, :, Plan.POSITION], + "path_std": plan_stds[:, np.newaxis, :, Plan.POSITION], + "path_prob": np.ones((batch_size, 1), dtype=np.float32), } @@ -199,31 +239,3 @@ def create_teacher( ) -> TeacherModel: """Factory function to create teacher model.""" return TeacherModel(models_dir=models_dir, fp16=fp16) - - -# Utility to parse model outputs -def parse_policy_outputs( - raw_outputs: np.ndarray, - num_hypotheses: int = 5, - horizon: int = 33, -) -> dict[str, np.ndarray]: - """Parse raw policy outputs into structured predictions. - - The policy outputs 1000 values encoding: - - Path predictions for multiple hypotheses - - Lane lines - - Road edges - - Lead vehicle info - - This is a placeholder - actual parsing requires reverse-engineering - the exact output format from openpilot's model metadata. - """ - # TODO: Implement proper output parsing based on model metadata - # For now, return a simplified structure - batch_size = raw_outputs.shape[0] - - return { - "path_mean": np.zeros((batch_size, num_hypotheses, horizon, 2)), - "path_std": np.ones((batch_size, num_hypotheses, horizon, 2)), - "path_prob": np.ones((batch_size, num_hypotheses)) / num_hypotheses, - } diff --git a/openpilot/tools/dgx/training/train.py b/openpilot/tools/dgx/training/train.py index d6cc2094b4ba9d..de96ca532b3dae 100755 --- a/openpilot/tools/dgx/training/train.py +++ b/openpilot/tools/dgx/training/train.py @@ -17,7 +17,11 @@ import torch.optim as optim # type: ignore[import-not-found] from torch.utils.data import DataLoader # type: ignore[import-not-found] +import numpy as np + # Local imports +from openpilot.selfdrive.modeld.constants import ModelConstants, Plan +from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict from openpilot.tools.dgx.training.dora import apply_dora_to_model, count_parameters, get_dora_parameters from openpilot.tools.dgx.training.losses import CombinedTrainingLoss @@ -111,14 +115,89 @@ def create_route_dataloader(data_path: str, batch_size: int): return create_dataloader(dataset, batch_size=batch_size) +class DummyTeacher: + """Zero-label stand-in for --dry-run, shaped like TeacherModel outputs.""" + + def __init__(self, metadata: dict): + self.input_shapes = metadata["input_shapes"] + self.output_slices = metadata["output_slices"] + self.output_len = next(iter(metadata["output_shapes"].values()))[-1] + + def generate_labels(self, img, big_img, desire, traffic_convention, **kwargs) -> dict[str, np.ndarray]: + b = img.shape[0] + return { + "features": np.zeros((b, ModelConstants.FEATURE_LEN), dtype=np.float32), + "raw_outputs": np.zeros((b, self.output_len), dtype=np.float32), + "path_mean": np.zeros((b, 1, ModelConstants.IDX_N, 3), dtype=np.float32), + "path_std": np.ones((b, 1, ModelConstants.IDX_N, 3), dtype=np.float32), + "path_prob": np.ones((b, 1), dtype=np.float32), + } + + +def student_forward( + student: nn.Module, + img: torch.Tensor, + big_img: torch.Tensor, + desire: torch.Tensor, + traffic_convention: torch.Tensor, + input_shapes: dict[str, tuple[int, ...]], + device: torch.device, +) -> torch.Tensor: + """Run the supercombo student and return flat (batch, N) outputs. + + The ONNX graph has fixed batch-1 shapes, so samples run one at a time. + Inputs are passed positionally in graph-input order; desire goes into the + last desire_pulse step (as modeld does), recurrent/action inputs are zero. + """ + outs = [] + for i in range(img.shape[0]): + inputs = [] + for name, shape in input_shapes.items(): + if name == "img": + t = img[i : i + 1].float() + elif name == "big_img": + t = big_img[i : i + 1].float() + elif name == "desire_pulse": + t = torch.zeros(shape, dtype=torch.float32, device=device) + t[0, -1, :] = desire[i].float() + elif name == "traffic_convention": + t = traffic_convention[i : i + 1].float() + else: # features_buffer, action_t: cold start + t = torch.zeros(shape, dtype=torch.float32, device=device) + inputs.append(t) + out = student(*inputs) + if isinstance(out, (list, tuple)): + out = out[0] + outs.append(out.reshape(1, -1)) + return torch.cat(outs, dim=0) + + +def extract_path_distribution(flat: torch.Tensor, output_slices: dict[str, slice]) -> tuple[torch.Tensor, torch.Tensor]: + """Differentiably extract the plan position MDN from flat model outputs. + + Mirrors Parser.parse_mdn for the single-hypothesis plan head: first half of + the slice is the mean, second half log-std. Returns path position mean/std + shaped (batch, 1, IDX_N, 3) for the distillation loss. + """ + plan_raw = flat[:, output_slices["plan"]] + b = plan_raw.shape[0] + n = plan_raw.shape[1] // 2 + shape = (b, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + mean = plan_raw[:, :n].reshape(shape) + std = torch.exp(plan_raw[:, n : 2 * n].clamp(max=11)).reshape(shape) + return mean[:, None, :, Plan.POSITION], std[:, None, :, Plan.POSITION] + + def train_epoch( student: nn.Module, - teacher, # TeacherModel (TensorRT) + teacher, # TeacherModel (TensorRT) or DummyTeacher dataloader: DataLoader, optimizer: optim.Optimizer, criterion: nn.Module, device: torch.device, epoch: int, + input_shapes: dict[str, tuple[int, ...]], + output_slices: dict[str, slice], log_interval: int = 10, ) -> dict[str, float]: """Train for one epoch.""" @@ -129,12 +208,19 @@ def train_epoch( start_time = time.perf_counter() for batch_idx, batch in enumerate(dataloader): - # Move to device - img = batch["img"].to(device) - big_img = batch["big_img"].to(device) + # Move to device (route dataloader emits road_frame/wide_frame) + img = batch.get("img", batch.get("road_frame")).to(device) + big_img = batch.get("big_img", batch.get("wide_frame")).to(device) desire = batch["desire"].to(device) traffic = batch["traffic_convention"].to(device) + # The model wants two temporally stacked 6-channel frames (12 channels); + # single-frame samples are duplicated as a cold-start approximation + if img.shape[1] == 6: + img = torch.cat([img, img], dim=1) + if big_img.shape[1] == 6: + big_img = torch.cat([big_img, big_img], dim=1) + # Generate teacher labels (no grad, uses TensorRT) with torch.no_grad(): # Convert to numpy for TensorRT @@ -145,15 +231,13 @@ def train_epoch( traffic_convention=traffic.cpu().numpy(), ) - # Student forward pass - # TODO: Implement proper forward pass based on model architecture - student_pred = student(img.float(), big_img.float()) + # Student forward pass: flat (batch, N) supercombo outputs + student_flat = student_forward(student, img, big_img, desire, traffic, input_shapes, device) + path_mean, path_std = extract_path_distribution(student_flat, output_slices) - # Compute loss - # TODO: Format predictions properly loss_dict = criterion( - student_pred={"raw": student_pred}, - teacher_pred={"raw": torch.from_numpy(teacher_labels["raw_outputs"]).to(device)}, + student_pred={"path_mean": path_mean, "path_std": path_std}, + teacher_pred={k: torch.from_numpy(teacher_labels[k]).float().to(device) for k in ("path_mean", "path_std", "path_prob")}, ) loss = loss_dict["total"] @@ -180,7 +264,7 @@ def train_epoch( def main(): parser = argparse.ArgumentParser(description="DoRA fine-tuning for openpilot") parser.add_argument("--data", type=str, default=None, help="Training data: path, 'ci' for CI segments, or 'comma_car_segments'") - parser.add_argument("--model", type=str, default="openpilot/selfdrive/modeld/models/driving_policy.onnx") + parser.add_argument("--model", type=str, default="openpilot/selfdrive/modeld/models/driving_supercombo.onnx") parser.add_argument("--epochs", type=int, default=10, help="Number of epochs") parser.add_argument("--batch-size", type=int, default=32, help="Batch size") parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate") @@ -230,11 +314,14 @@ def main(): # Setup loss criterion = CombinedTrainingLoss(path_weight=1.0, feature_weight=0.1) + # Model I/O specs from the ONNX-embedded metadata (input_shapes, output_slices) + metadata = make_metadata_dict(args.model) + # Setup teacher (TensorRT) print("\nLoading teacher model...") if args.dry_run: print("Dry run - using dummy teacher") - teacher = None # Would use dummy labels + teacher = DummyTeacher(metadata) else: from openpilot.tools.dgx.training.teacher import create_teacher @@ -264,6 +351,8 @@ def main(): criterion=criterion, device=device, epoch=epoch, + input_shapes=metadata["input_shapes"], + output_slices=metadata["output_slices"], ) print(f"\nEpoch {epoch} complete:") From f5fabf8741410b4766bd980c457da270d4ecdcb0 Mon Sep 17 00:00:00 2001 From: Evan Montgomery-Recht Date: Fri, 3 Jul 2026 18:40:03 -0400 Subject: [PATCH 2/5] fork: fix OnnxRunner import for pinned tinygrad tinygrad moved OnnxRunner to tinygrad.nn.onnx (matching upstream's compile_modeld.py); tinygrad.frontend.onnx no longer exists. Found running the benchmark on the DGX Spark. Co-Authored-By: Claude Fable 5 --- openpilot/tools/dgx/benchmark_inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openpilot/tools/dgx/benchmark_inference.py b/openpilot/tools/dgx/benchmark_inference.py index bfcc5de512b871..cc38520eb71c4b 100755 --- a/openpilot/tools/dgx/benchmark_inference.py +++ b/openpilot/tools/dgx/benchmark_inference.py @@ -45,7 +45,7 @@ def main(): os.environ["BEAM"] = str(args.beam) from tinygrad import Device, Tensor - from tinygrad.frontend.onnx import OnnxRunner # type: ignore[import-not-found] + from tinygrad.nn.onnx import OnnxRunner # type: ignore[import-not-found] Device.DEFAULT = "CUDA" print(f"Device: {Device.DEFAULT}") From 0bcc5ebee2a6274c1f538fae19546642ee7479b5 Mon Sep 17 00:00:00 2001 From: Evan Montgomery-Recht Date: Fri, 3 Jul 2026 18:42:47 -0400 Subject: [PATCH 3/5] fork: set tinygrad device via DEV env, Device.DEFAULT is read-only now Found running the benchmark on the DGX Spark. Co-Authored-By: Claude Fable 5 --- openpilot/tools/dgx/benchmark_inference.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/openpilot/tools/dgx/benchmark_inference.py b/openpilot/tools/dgx/benchmark_inference.py index cc38520eb71c4b..8e36dfd98576d4 100755 --- a/openpilot/tools/dgx/benchmark_inference.py +++ b/openpilot/tools/dgx/benchmark_inference.py @@ -43,11 +43,12 @@ def main(): if args.beam > 0: os.environ["BEAM"] = str(args.beam) + # device must be chosen via env before tinygrad import (Device.DEFAULT is read-only now) + os.environ.setdefault("DEV", "CUDA") from tinygrad import Device, Tensor from tinygrad.nn.onnx import OnnxRunner # type: ignore[import-not-found] - Device.DEFAULT = "CUDA" print(f"Device: {Device.DEFAULT}") print(f"BEAM: {os.environ.get('BEAM', 'disabled')}") print("=" * 60) From 126e1a8ad091d2358f6063b4ca6b30a3ad6bef58 Mon Sep 17 00:00:00 2001 From: Evan Montgomery-Recht Date: Fri, 3 Jul 2026 18:51:16 -0400 Subject: [PATCH 4/5] fork: record DGX Spark validation results, document DoRA converter gap Measured on GB10 (TensorRT 10.14, driver 580.159, 2026-07-03): - driving_supercombo: 1.07ms (937 FPS), engine builds in 23s - dmonitoring: 0.26ms (3,891 FPS); combined pipeline 1.32ms (755 FPS) - teacher generate_labels validated end-to-end on GPU (all parsed outputs, recurrent feature feedback roundtrip) - tinygrad sees CUDA (SConscript build-time selection would pick it); tinygrad baseline still needs clang installed on the box Known limitation, validated in practice: neither onnx2pytorch nor onnx2torch can convert the opset-20 supercombo to PyTorch (Cast v19, Gelu v20, Reshape allowzero, axes-as-input Reduce* v18 all unsupported), so the DoRA student path fails at model load. load_student_model now raises an actionable error instead of a KeyError four layers deep; README/TRAINING document alternatives (tinygrad fine-tuning or pinning the last split-model release). Co-Authored-By: Claude Fable 5 --- .../tools/dgx/HARDWARE_TEST_CHECKLIST.md | 10 ++--- openpilot/tools/dgx/README.md | 41 ++++++++++--------- openpilot/tools/dgx/TRAINING.md | 11 ++++- openpilot/tools/dgx/training/train.py | 16 +++++++- 4 files changed, 51 insertions(+), 27 deletions(-) diff --git a/openpilot/tools/dgx/HARDWARE_TEST_CHECKLIST.md b/openpilot/tools/dgx/HARDWARE_TEST_CHECKLIST.md index 2563b62d456e09..4885d3b4135ea1 100644 --- a/openpilot/tools/dgx/HARDWARE_TEST_CHECKLIST.md +++ b/openpilot/tools/dgx/HARDWARE_TEST_CHECKLIST.md @@ -61,11 +61,11 @@ print(f'EON: {HARDWARE.get_device_type()}') ```bash python openpilot/tools/dgx/benchmark_tensorrt.py --runs 50 --warmup 10 ``` -Expected results (DGX Spark GB10): -- [ ] driving_supercombo: < 2.5ms (target TBD on hardware; placeholder is the sum of the old split-model targets) -- [ ] dmonitoring: < 1ms (>1000 FPS) -- [ ] Combined: < 3.5ms (target TBD on hardware) -- [ ] No CUDA errors or warnings +Expected results (DGX Spark GB10; measured 2026-07-03, TensorRT 10.14): +- [x] driving_supercombo: < 2.5ms (measured 1.07ms, 937 FPS) +- [x] dmonitoring: < 1ms (measured 0.26ms, 3,891 FPS) +- [x] Combined: < 3.5ms (measured 1.32ms, 755 FPS) +- [x] No CUDA errors (two benign FP16 clamp precision warnings during engine build) ### 2.2 tinygrad CUDA Benchmark ```bash diff --git a/openpilot/tools/dgx/README.md b/openpilot/tools/dgx/README.md index ba05782eebf8f9..759fef3e817415 100644 --- a/openpilot/tools/dgx/README.md +++ b/openpilot/tools/dgx/README.md @@ -214,33 +214,29 @@ Options: - `--warmup N`: Number of warmup runs (default: 5) - `--beam N`: BEAM optimization level (default: 0, disabled) -### Benchmark Results (2026-01, historical) +### Benchmark Results Tested on DGX Spark (GB10, Blackwell, compute 12.1). -**Note**: these numbers were measured against the pre-July-2026 split models -(`driving_vision.onnx` / `driving_policy.onnx`), which have been replaced upstream -by the combined `driving_supercombo.onnx`. The rows are kept as historical -reference; supercombo numbers are pending re-validation on DGX hardware. - #### TensorRT FP16 (Recommended) -Historical measurements (pre-July-2026 split models): +Measured 2026-07-03 (TensorRT 10.14, driver 580.159, combined model): -| Model | Inference | FPS | vs tinygrad | -|-------|-----------|-----|-------------| -| driving_policy (historical) | 0.09ms | 11,355 | 659x faster | -| driving_vision (historical) | 0.85ms | 1,181 | 432x faster | -| dmonitoring | 0.28ms | 3,584 | 1,175x faster | -| **Combined (historical)** | **1.21ms** | **824** | **620x faster** | -| driving_supercombo | pending re-validation on DGX hardware | — | — | +| Model | Inference | FPS | +|-------|-----------|-----| +| driving_supercombo | 1.07ms ± 0.10ms | 937 | +| dmonitoring | 0.26ms ± 0.03ms | 3,891 | +| **Combined pipeline** | **1.32ms** | **755** | -**TensorRT was 41x faster than comma 3X** (1.2ms vs ~50ms) in the historical -split-model benchmark. +Engine build time: 23s (supercombo), 8s (dmonitoring). **TensorRT is ~38x +faster than the comma 3X real-time budget** (1.3ms vs ~50ms). + +Historical split-model measurements (2026-01, pre-restructure): driving_policy +0.09ms / driving_vision 0.85ms / combined 1.21ms (824 FPS, 620x over tinygrad). #### tinygrad CUDA (Not Optimized for Blackwell) -Historical measurements (pre-July-2026 split models): +Historical measurements (2026-01, pre-July-2026 split models): | Model | Inference | FPS | |-------|-----------|-----| @@ -248,7 +244,7 @@ Historical measurements (pre-July-2026 split models): | driving_vision (historical) | 366ms | 2.7 | | dmonitoring | 328ms | 3.0 | | **Combined (historical)** | **514ms** | **1.9** | -| driving_supercombo | pending re-validation on DGX hardware | — | +| driving_supercombo | not yet measured (tinygrad's CPU-side compiler needs `clang` installed) | tinygrad's CUDA backend is not yet optimized for Blackwell architecture. Use TensorRT for production-level performance. @@ -261,9 +257,16 @@ The DGX Spark is ideal for fine-tuning openpilot models using DoRA (Weight-Decom - **Parameter Efficient**: Only ~2-5% of parameters are trained - **Preserves Base Model**: Original weights are frozen, preventing catastrophic forgetting -- **Fast Training**: With TensorRT teacher, pseudo-label generation ran at 800+ FPS (historical split-model measurement; supercombo pending re-validation) +- **Fast Training**: With TensorRT teacher, pseudo-label generation runs at 937 FPS (supercombo, measured on GB10 2026-07) - **Easy Deployment**: DoRA weights can be merged back into the base model for inference +> **Known limitation (2026-07)**: the student-model path is currently blocked — +> neither onnx2pytorch nor onnx2torch can convert the opset-20 +> `driving_supercombo.onnx` to PyTorch (missing Cast v19, Gelu v20, Reshape +> `allowzero`, and axes-as-input Reduce* v18 support). The TensorRT **teacher** +> works fully. Until a converter catches up: fine-tune in tinygrad (the model +> already runs there), or pin the last split-model release for torch experiments. + ### Quick Start Training ```bash diff --git a/openpilot/tools/dgx/TRAINING.md b/openpilot/tools/dgx/TRAINING.md index 53eb09ffab5d64..83139ad7308f8b 100644 --- a/openpilot/tools/dgx/TRAINING.md +++ b/openpilot/tools/dgx/TRAINING.md @@ -63,7 +63,7 @@ July 2026 upstream restructure there is no separate vision stage: the combined `driving_supercombo.onnx` does vision + policy in a single forward pass. ```python -# Teacher generates targets (800+ FPS measured historically on the split models) +# Teacher generates targets (937 FPS measured on GB10, TensorRT 10.14, 2026-07) teacher = TensorRTEngine("driving_supercombo.onnx") # Model metadata (input_shapes, output_shapes, output_slices, model_checkpoint) @@ -99,6 +99,13 @@ desire_state, pad) and parsed with `Parser().parse_outputs()` from ### 3. DoRA Fine-Tuning +> **Known limitation (validated on DGX Spark, 2026-07)**: loading the student +> currently fails — neither onnx2pytorch nor onnx2torch can convert the +> opset-20 `driving_supercombo.onnx` to PyTorch (missing Cast v19, Gelu v20, +> Reshape `allowzero`, and axes-as-input Reduce* v18). The TensorRT teacher +> below works fully. Alternatives: fine-tune in tinygrad, or pin the last +> split-model release for torch experiments. + **DoRA** (Weight-Decomposed Low-Rank Adaptation): - Decomposes weights into magnitude and direction - Only trains low-rank direction updates @@ -260,7 +267,7 @@ openpilot/tools/dgx/ DGX Spark advantages: - 128GB unified memory = large batch sizes -- TensorRT teacher at 800+ FPS (historical split-model measurement) = fast pseudo-label generation +- TensorRT teacher at 937 FPS (supercombo, measured on GB10 2026-07) = fast pseudo-label generation - Single device = no distributed training complexity ## Usage diff --git a/openpilot/tools/dgx/training/train.py b/openpilot/tools/dgx/training/train.py index de96ca532b3dae..f29b594d7b7cae 100755 --- a/openpilot/tools/dgx/training/train.py +++ b/openpilot/tools/dgx/training/train.py @@ -30,6 +30,13 @@ def load_student_model(onnx_path: str, device: torch.device) -> nn.Module: """Load student model from ONNX and convert to PyTorch. Uses onnx2pytorch for conversion, then wraps for training. + + KNOWN LIMITATION (validated on DGX Spark, 2026-07): neither onnx2pytorch + (0.5.3) nor onnx2torch can convert the opset-20 driving_supercombo.onnx — + they miss Cast v19, Gelu v20, Reshape allowzero=1, and the v18 + axes-as-input Reduce* ops. Options until a converter catches up: + fine-tune in tinygrad (the model already runs there), or pin the last + split-model release (pre July 2026) for torch-based experiments. """ try: import onnx @@ -39,7 +46,14 @@ def load_student_model(onnx_path: str, device: torch.device) -> nn.Module: print(f"Loading student model from {onnx_path}...") onnx_model = onnx.load(onnx_path) - pytorch_model = ConvertModel(onnx_model) + try: + pytorch_model = ConvertModel(onnx_model) + except (KeyError, NotImplementedError) as e: + opsets = [(o.domain or "ai.onnx", o.version) for o in onnx_model.opset_import] + msg = f"onnx2pytorch cannot convert {onnx_path} (opsets {opsets}): {e!r}." + msg += " The combined driving_supercombo.onnx uses opset 20, which current ONNX->PyTorch converters do not support" + msg += " (Cast v19, Gelu v20, Reshape allowzero, Reduce* v18). See the load_student_model docstring for options." + raise RuntimeError(msg) from e pytorch_model = pytorch_model.to(device) return pytorch_model From f2eab71ae80d88170dcbaa9359dd3c7c39d62a0c Mon Sep 17 00:00:00 2001 From: Evan Montgomery-Recht Date: Fri, 3 Jul 2026 21:27:36 -0400 Subject: [PATCH 5/5] fork: appease codespell (unparseable -> unparsable) Co-Authored-By: Claude Fable 5 --- openpilot/tools/dgx/model_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openpilot/tools/dgx/model_runner.py b/openpilot/tools/dgx/model_runner.py index 812b8461cdd27c..71acf4529f9c79 100644 --- a/openpilot/tools/dgx/model_runner.py +++ b/openpilot/tools/dgx/model_runner.py @@ -340,7 +340,7 @@ def _create_dummy_inputs(self) -> dict[str, np.ndarray]: from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict self.input_shapes = make_metadata_dict(self.model_path)["input_shapes"] - except Exception: # no embedded metadata / unparseable — fall through + except Exception: # no embedded metadata / unparsable — fall through pass if not self.input_shapes: