diff --git a/Model/tests/test_workflow_training_lifecycle.py b/Model/tests/test_workflow_training_lifecycle.py index c874aad3a..5012b7d15 100644 --- a/Model/tests/test_workflow_training_lifecycle.py +++ b/Model/tests/test_workflow_training_lifecycle.py @@ -1688,3 +1688,46 @@ def test_resume_load_keeps_rng_tensors_on_cpu(): keywords = {item.arg: item.value for item in resume_load.keywords} assert ast.literal_eval(keywords["map_location"]) == "cpu" assert ast.literal_eval(keywords["weights_only"]) is False + + +def test_camera_bev_grid_override_changes_only_the_grid(): + """The 6 GB workaround must not silently change what the BEV covers. + + Shrinking the grid is safe only if pc_range is left alone: the BEV then + covers the same ground area with coarser cells, and the map BEV stays + aligned. If pc_range moved too, the model would be looking at a smaller + patch of the world and the metric would not be comparable. + """ + from navigation.geometry import DEFAULT_NAVIGATION_GEOMETRY + + base = DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs() + resized = workflows._camera_bev_kwargs_with_grid(base, 64) + + assert resized["bev_h"] == 64 + assert resized["bev_w"] == 64 + changed = {k for k in base if base[k] != resized.get(k)} + assert changed == {"bev_h", "bev_w"}, ( + f"resizing the grid must not touch anything else, but changed {changed}" + ) + assert base["bev_h"] != 64, "fixture would be vacuous if the default were 64" + # The caller must not have its own dict mutated underneath it. + assert base["bev_h"] == DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()["bev_h"] + + +@pytest.mark.parametrize("bad", [0, -1, -256]) +def test_camera_bev_grid_override_rejects_non_positive(bad): + from navigation.geometry import DEFAULT_NAVIGATION_GEOMETRY + + base = DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs() + with pytest.raises(ValueError, match="camera_bev_size must be positive"): + workflows._camera_bev_kwargs_with_grid(base, bad) + + +def test_camera_bev_grid_defaults_to_the_geometry(): + """With the parameter unset the run must be byte-identical to before.""" + source = inspect.getsource(workflows.train_il.task_function) + assert "camera_bev_size: Optional[int] = None" in source + assert "if camera_bev_size is not None:" in source, ( + "the override must be opt-in; an unconditional call would change the " + "default grid for every existing run" + ) diff --git a/Platform/pipelines/workflows.py b/Platform/pipelines/workflows.py index 48b239ed7..f2de5531f 100644 --- a/Platform/pipelines/workflows.py +++ b/Platform/pipelines/workflows.py @@ -403,6 +403,27 @@ def _training_num_views_from_manifests( return max(dataset_views.values()) +def _camera_bev_kwargs_with_grid( + base: dict, + size: int, +) -> dict: + """Return ``base`` with the camera BEV grid resized to ``size`` per side. + + Only ``bev_h``/``bev_w`` change. ``pc_range`` is deliberately untouched, so + the BEV still covers the ground area the navigation geometry defines and the + map BEV stays aligned with it -- the grid is coarser, not smaller. + + The KITScenes geometry pins the grid at 256, which allocates 65,536 BEV + queries and does not fit in 6 GB of VRAM. + """ + if size < 1: + raise ValueError(f"camera_bev_size must be positive, got {size}") + resized = dict(base) + resized["bev_h"] = size + resized["bev_w"] = size + return resized + + def _training_source_revision( manifests: dict[str, dict], *, @@ -3111,6 +3132,11 @@ def train_il( # GPU step. Effective parallelism is capped by shard count, so scale needs more # (smaller) shards too. num_workers: int = 0, + # Camera BEV grid, in cells per side. None keeps the value the KITScenes + # geometry defines (256), which allocates 65,536 queries and does not fit in + # 6 GB of VRAM. Only the grid changes: pc_range stays as the geometry + # defines it, so the BEV covers the same ground area with coarser cells. + camera_bev_size: Optional[int] = None, resume_from: Optional[FlyteFile] = None, early_stopping_patience: int = 5, allow_resume_policy_transition: bool = False, @@ -3387,6 +3413,16 @@ def train_il( view_fusion_kwargs = ( DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs() ) + if camera_bev_size is not None: + view_fusion_kwargs = _camera_bev_kwargs_with_grid( + view_fusion_kwargs, camera_bev_size + ) + print( + f"Camera BEV grid overridden: {camera_bev_size}x{camera_bev_size} " + f"({camera_bev_size ** 2} cells; geometry default is " + f"{DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()['bev_h']}x" + f"{DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()['bev_w']})" + ) from Platform.pipelines.training_checkpoint import stable_digest @@ -8698,6 +8734,7 @@ def wf_train_il( val_fraction: float = 0.1, validation_scope: str = "full", num_workers: int = 0, + camera_bev_size: Optional[int] = None, resume_from: Optional[FlyteFile] = None, early_stopping_patience: int = 5, allow_resume_policy_transition: bool = False, @@ -8733,7 +8770,8 @@ def wf_train_il( enable_reasoning=enable_reasoning, reasoning_mode=reasoning_mode, enable_world_model=enable_world_model, val_fraction=val_fraction, validation_scope=validation_scope, - num_workers=num_workers, resume_from=resume_from, + num_workers=num_workers, camera_bev_size=camera_bev_size, + resume_from=resume_from, early_stopping_patience=early_stopping_patience, allow_resume_policy_transition=allow_resume_policy_transition) return evaluate_il_policy(checkpoint=out.checkpoint, shards=shards, dataset=dataset,