mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-08-05 15:35:40 +08:00
simplify everything.
This commit is contained in:
nayan
@@ -1,128 +0,0 @@
|
||||
"""Infer MHP (mixture-density-hypothesis) parameter values from output slice sizes.
|
||||
|
||||
The legacy supercombo encoded its outputs using fixed values for
|
||||
``PLAN_MHP_N``, ``PLAN_MHP_SELECTION``, ``LEAD_MHP_N``, and ``LEAD_MHP_SELECTION``.
|
||||
Newer supercombo architectures emit differently-sized slices for the same heads,
|
||||
yet the parser downstream still expects to be told the layout.
|
||||
|
||||
This module figures those numbers out at runtime from the model's output slice,
|
||||
so the existing ``parse_mdn`` path can handle any supercombo flavor without code
|
||||
changes for each variant. ``infer_mhp`` always tries the legacy values first so
|
||||
existing compiled pkls behave identically (full backwards compatibility).
|
||||
|
||||
Schema reminder (``parse_mdn`` packs hypotheses contiguously in the order
|
||||
``in_N × (mu | std | weights)`` along the channel axis):
|
||||
|
||||
slice_size == in_N × (2·n_values + out_N)
|
||||
|
||||
where ``n_values`` is the per-hypothesis value width (``IDX_N × PLAN_WIDTH``
|
||||
for plan, ``LEAD_TRAJ_LEN × LEAD_WIDTH`` for lead, etc.).
|
||||
"""
|
||||
|
||||
def infer_mhp(
|
||||
slice_size: int,
|
||||
n_values: int,
|
||||
legacy_in_n: int,
|
||||
legacy_out_n: int,
|
||||
max_in_n: int = 16,
|
||||
) -> tuple[int, int]:
|
||||
"""Infer ``(in_N, out_N)`` for an MDN-encoded output slice.
|
||||
|
||||
Tries values in this order, returning on the first match:
|
||||
|
||||
1. ``(legacy_in_n, legacy_out_n)`` exactly — preserves exact backwards
|
||||
compatibility for existing supercombo pkls.
|
||||
2. ``out_N ∈ (0, 1, 3)`` (no weights, single weight, three-way selection) at
|
||||
increasing in_N — covers the common architectural patterns.
|
||||
3. Brute-force any valid ``out_N`` that divides ``slice_size``.
|
||||
|
||||
Args:
|
||||
slice_size: Number of floats in the output slice (typically
|
||||
``slices[name].stop - slices[name].start``).
|
||||
n_values: Per-hypothesis mu/std width (e.g. ``IDX_N × PLAN_WIDTH`` for plan).
|
||||
legacy_in_n: The legacy in_N value (highest priority for backwards compat).
|
||||
legacy_out_n: The legacy out_N value.
|
||||
max_in_n: Upper bound on accepted hypothesis counts (filters silly parses).
|
||||
|
||||
Returns:
|
||||
``(in_N, out_N)``. If nothing fits the formulas, returns ``(1, 0)`` —
|
||||
single hypothesis with no weights — which is the gentlest fallback.
|
||||
"""
|
||||
if slice_size <= 0 or n_values <= 0:
|
||||
return 1, 0
|
||||
|
||||
# Priority 1: exact legacy match (BC-preserving).
|
||||
per_hyp_legacy = 2 * n_values + legacy_out_n
|
||||
if per_hyp_legacy > 0 and legacy_in_n * per_hyp_legacy == slice_size:
|
||||
return legacy_in_n, legacy_out_n
|
||||
|
||||
# Priority 2: common weight layouts across supercombo variants.
|
||||
for out_n in (0, 1, 3):
|
||||
per_hyp = 2 * n_values + out_n
|
||||
if per_hyp <= 0:
|
||||
continue
|
||||
if slice_size % per_hyp == 0:
|
||||
in_n = slice_size // per_hyp
|
||||
if 1 <= in_n <= max_in_n:
|
||||
return in_n, out_n
|
||||
|
||||
# Priority 3: brute-force any divisor that yields a sensible in_N.
|
||||
# Bound out_n by max_in_n to keep the search tiny (3 hypotheses
|
||||
# of weights is already a lot).
|
||||
for out_n in range(0, max_in_n + 1):
|
||||
per_hyp = 2 * n_values + out_n
|
||||
if per_hyp <= 0:
|
||||
continue
|
||||
if slice_size % per_hyp == 0:
|
||||
in_n = slice_size // per_hyp
|
||||
if 1 <= in_n <= max_in_n:
|
||||
return in_n, out_n
|
||||
|
||||
# Last resort: best-effort single hypothesis with no weights.
|
||||
return 1, 0
|
||||
|
||||
|
||||
def slice_size(sl) -> int:
|
||||
"""Return the float-width of a ``slice``/``None`` from ``output_slices``.
|
||||
|
||||
``None`` and ``slice(None, None, None)`` are treated as 0 (output absent).
|
||||
Negative-end slices (ONNX-style "from end") aren't supported by this helper
|
||||
because parser expects single-tensor packed outputs.
|
||||
"""
|
||||
if sl is None:
|
||||
return 0
|
||||
start = 0 if sl.start is None else sl.start
|
||||
stop = sl.stop
|
||||
if stop is None or stop < 0:
|
||||
return 0
|
||||
return max(0, stop - start)
|
||||
|
||||
|
||||
def infer_mhp_for_outputs(
|
||||
output_slices: dict,
|
||||
constants,
|
||||
max_in_n: int = 16,
|
||||
) -> dict:
|
||||
"""Build a dict of MHP values keyed by head name from a model's output_slices.
|
||||
|
||||
Reads ``output_slices['plan']`` and ``output_slices['lead']`` (if present)
|
||||
and infers their ``in_N``/``out_N``. Other heads aren't MDN-encoded the same
|
||||
way and remain driven by the ``constants`` module.
|
||||
"""
|
||||
config: dict[str, int] = {}
|
||||
|
||||
plan_size = slice_size(output_slices.get('plan'))
|
||||
if plan_size > 0:
|
||||
n = constants.IDX_N * constants.PLAN_WIDTH
|
||||
in_n, out_n = infer_mhp(plan_size, n, constants.PLAN_MHP_N, constants.PLAN_MHP_SELECTION, max_in_n)
|
||||
config['plan_mhp_n'] = in_n
|
||||
config['plan_mhp_selection'] = out_n
|
||||
|
||||
lead_size = slice_size(output_slices.get('lead'))
|
||||
if lead_size > 0:
|
||||
n = constants.LEAD_TRAJ_LEN * constants.LEAD_WIDTH
|
||||
in_n, out_n = infer_mhp(lead_size, n, constants.LEAD_MHP_N, constants.LEAD_MHP_SELECTION, max_in_n)
|
||||
config['lead_mhp_n'] = in_n
|
||||
config['lead_mhp_selection'] = out_n
|
||||
|
||||
return config
|
||||
@@ -156,25 +156,15 @@ class ModelState(ModelStateBase):
|
||||
from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
|
||||
self.constants = ModelConstants()
|
||||
|
||||
# Derive the parser's per-head MHP values from the appropriate output
|
||||
# slices. Legacy pkls fall back to the constants values (Priority 1 in
|
||||
# ``mhp_inference.infer_mhp``), so this is fully backwards compatible.
|
||||
# Supercombo pkls carry plan/lead in the vision tensor; split/multi-policy
|
||||
# pkls carry them on the policy tensor (and even there we use the first
|
||||
# policy's slices -- the existing code only tracks one ``policy_output_slices``).
|
||||
from openpilot.sunnypilot.modeld_v2.mhp_inference import infer_mhp_for_outputs
|
||||
|
||||
if self._combined_model_type == 'supercombo':
|
||||
mhp_config = infer_mhp_for_outputs(self.vision_output_slices, self.constants)
|
||||
else:
|
||||
mhp_config = infer_mhp_for_outputs(self.policy_output_slices, self.constants)
|
||||
|
||||
# Combined parsers auto-detect ``(in_N, out_N)`` for plan/lead from the raw
|
||||
# slice size, so they transparently support both legacy and newer
|
||||
# supercombo ONNXes without any per-model configuration.
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser
|
||||
if self._combined_model_type != 'supercombo':
|
||||
self.parser = SplitParser(mhp_config=mhp_config)
|
||||
self.parser = SplitParser()
|
||||
else:
|
||||
self.parser = CombinedParser(mhp_config=mhp_config)
|
||||
self.parser = CombinedParser()
|
||||
|
||||
self.prev_desire = np.zeros(self.constants.DESIRE_LEN, dtype=np.float32)
|
||||
self.full_frames: dict = {}
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
import numpy as np
|
||||
from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
|
||||
|
||||
|
||||
def safe_exp(x, out=None):
|
||||
# -11 is around 10**14, more causes float16 overflow
|
||||
return np.exp(np.clip(x, -np.inf, 11), out=out)
|
||||
|
||||
|
||||
def sigmoid(x):
|
||||
return 1. / (1. + safe_exp(-x))
|
||||
|
||||
|
||||
def softmax(x, axis=-1):
|
||||
x -= np.max(x, axis=axis, keepdims=True)
|
||||
if x.dtype == np.float32 or x.dtype == np.float64:
|
||||
@@ -17,17 +20,34 @@ def softmax(x, axis=-1):
|
||||
x /= np.sum(x, axis=axis, keepdims=True)
|
||||
return x
|
||||
|
||||
class Parser:
|
||||
def __init__(self, ignore_missing=False, mhp_config=None):
|
||||
self.ignore_missing = ignore_missing
|
||||
# Optional MHP overrides keyed by head: 'plan_mhp_n', 'plan_mhp_selection',
|
||||
# 'lead_mhp_n', 'lead_mhp_selection'. ``None`` (or missing keys) keeps the
|
||||
# legacy ``ModelConstants`` values so existing models behave identically.
|
||||
self.mhp = mhp_config or {}
|
||||
|
||||
def _mhp(self, head, default_in, default_out):
|
||||
return (self.mhp.get(f'{head}_mhp_n', default_in),
|
||||
self.mhp.get(f'{head}_mhp_selection', default_out))
|
||||
def _infer_mhp(slice_size: int, prod_out_shape: int, max_in_n: int = 16, max_out_n: int = 6) -> tuple[int, int]:
|
||||
"""Derive ``(in_N, out_N)`` from a packed MDN slice.
|
||||
|
||||
Layout (combined supercombo): for each hypothesis we have ``mu``, ``std``,
|
||||
and an optional scalar ``weight`` block. So:
|
||||
|
||||
slice_size = in_N * (2 * prod_out_shape + out_N)
|
||||
|
||||
We scan small ``out_N`` values (no weights is most common in modern models,
|
||||
one or three weights in the legacy) and accept the first division that
|
||||
yields an integer ``in_N`` in ``[1, max_in_n]``. The candidates are spaced
|
||||
wide enough that there's no ambiguity for the bands we care about.
|
||||
"""
|
||||
for out_n in range(max_out_n + 1):
|
||||
per = 2 * prod_out_shape + out_n
|
||||
if per <= 0:
|
||||
continue
|
||||
if slice_size % per == 0:
|
||||
in_n = slice_size // per
|
||||
if 1 <= in_n <= max_in_n:
|
||||
return in_n, out_n
|
||||
return 1, 0 # single hypothesis, no weights — matches a non-MDN output
|
||||
|
||||
|
||||
class Parser:
|
||||
def __init__(self, ignore_missing=False):
|
||||
self.ignore_missing = ignore_missing
|
||||
|
||||
def check_missing(self, outs, name):
|
||||
if name not in outs and not self.ignore_missing:
|
||||
@@ -48,85 +68,92 @@ class Parser:
|
||||
raw = outs[name]
|
||||
outs[name] = sigmoid(raw)
|
||||
|
||||
def parse_mdn(self, name, outs, in_N=0, out_N=1, out_shape=None):
|
||||
def parse_mdn(self, name, outs, out_shape, in_N=0, out_N=0):
|
||||
"""Parse a packed MDN output. Pass ``in_N``/``out_N`` explicitly for the
|
||||
legacy layout; pass neither (defaults of 0) to auto-detect from the
|
||||
slice size, which transparently supports newer supercombos that drop the
|
||||
per-hypothesis weight block or change the hypothesis count."""
|
||||
if self.check_missing(outs, name):
|
||||
return
|
||||
raw = outs[name]
|
||||
raw = raw.reshape((raw.shape[0], max(in_N, 1), -1))
|
||||
|
||||
if in_N == 0 and out_N == 0:
|
||||
prod = int(np.prod(out_shape))
|
||||
in_N, out_N = _infer_mhp(raw.shape[1], prod)
|
||||
|
||||
raw = raw.reshape((raw.shape[0], in_N, -1))
|
||||
|
||||
n_values = (raw.shape[2] - out_N)//2
|
||||
pred_mu = raw[:,:,:n_values]
|
||||
pred_std = safe_exp(raw[:,:,n_values: 2*n_values])
|
||||
|
||||
if in_N > 1:
|
||||
if out_N > 0:
|
||||
weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype)
|
||||
for i in range(out_N):
|
||||
weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1)
|
||||
if in_N > 1 and out_N > 0:
|
||||
weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype)
|
||||
for i in range(out_N):
|
||||
weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1)
|
||||
|
||||
if out_N == 1:
|
||||
for fidx in range(weights.shape[0]):
|
||||
idxs = np.argsort(weights[fidx][:,0])[::-1]
|
||||
weights[fidx] = weights[fidx][idxs]
|
||||
pred_mu[fidx] = pred_mu[fidx][idxs]
|
||||
pred_std[fidx] = pred_std[fidx][idxs]
|
||||
assert out_shape is not None
|
||||
full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
|
||||
outs[name + '_weights'] = weights
|
||||
outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
|
||||
outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
|
||||
|
||||
pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
|
||||
pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
|
||||
if out_N == 1:
|
||||
for fidx in range(weights.shape[0]):
|
||||
for hidx in range(out_N):
|
||||
idxs = np.argsort(weights[fidx,:,hidx])[::-1]
|
||||
pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]]
|
||||
pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]]
|
||||
else:
|
||||
# MHP without weights: keep every hypothesis intact, surface them as
|
||||
# ``*_hypotheses`` outputs and use the same shape for the primary
|
||||
# output so downstream consumers can iterate over the full set.
|
||||
assert out_shape is not None
|
||||
full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
|
||||
outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
|
||||
outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
|
||||
pred_mu_final = pred_mu
|
||||
pred_std_final = pred_std
|
||||
idxs = np.argsort(weights[fidx][:,0])[::-1]
|
||||
weights[fidx] = weights[fidx][idxs]
|
||||
pred_mu[fidx] = pred_mu[fidx][idxs]
|
||||
pred_std[fidx] = pred_std[fidx][idxs]
|
||||
full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
|
||||
outs[name + '_weights'] = weights
|
||||
outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
|
||||
outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
|
||||
|
||||
pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
|
||||
pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
|
||||
for fidx in range(weights.shape[0]):
|
||||
for hidx in range(out_N):
|
||||
idxs = np.argsort(weights[fidx,:,hidx])[::-1]
|
||||
pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]]
|
||||
pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]]
|
||||
elif in_N > 1 and out_N == 0:
|
||||
# MHP without weights: keep every hypothesis intact, surface them as
|
||||
# ``*_hypotheses`` and propagate the full multi-hypothesis tensor.
|
||||
full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
|
||||
outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
|
||||
outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
|
||||
pred_mu_final = pred_mu
|
||||
pred_std_final = pred_std
|
||||
else:
|
||||
pred_mu_final = pred_mu
|
||||
pred_std_final = pred_std
|
||||
|
||||
# Final-shape selector: keep an extra hypothesis axis only when
|
||||
# multi-hypothesis data must survive — i.e. when out_N > 1 (legacy
|
||||
# multiple-selection) OR when in_N > 1 with no weights (newer MHP). For
|
||||
# single-hypothesis / collapsed cases (lane_lines, pose, etc.) drop the
|
||||
# extra axis so the consumer sees the historical ``(batch, *out_shape)``.
|
||||
if out_N > 1 or (in_N > 1 and out_N == 0):
|
||||
assert out_shape is not None
|
||||
n_selections = out_N if out_N > 1 else in_N
|
||||
final_shape = tuple([raw.shape[0], n_selections] + list(out_shape))
|
||||
else:
|
||||
assert out_shape is not None
|
||||
final_shape = tuple([raw.shape[0],] + list(out_shape))
|
||||
outs[name] = pred_mu_final.reshape(final_shape)
|
||||
outs[name + '_stds'] = pred_std_final.reshape(final_shape)
|
||||
|
||||
def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
|
||||
plan_in, plan_out = self._mhp('plan', ModelConstants.PLAN_MHP_N, ModelConstants.PLAN_MHP_SELECTION)
|
||||
lead_in, lead_out = self._mhp('lead', ModelConstants.LEAD_MHP_N, ModelConstants.LEAD_MHP_SELECTION)
|
||||
self.parse_mdn('plan', outs, in_N=plan_in, out_N=plan_out,
|
||||
out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH))
|
||||
self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH))
|
||||
self.parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_ROAD_EDGES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH))
|
||||
self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
|
||||
self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
|
||||
# Pass no explicit ``in_N``/``out_N`` for plan/lead — ``parse_mdn`` infers
|
||||
# them from the raw slice size, which naturally handles both the legacy
|
||||
# supercombo (4955 / 102) and newer variants (e.g. 990 / 144).
|
||||
self.parse_mdn('plan', outs, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH))
|
||||
self.parse_mdn('lane_lines', outs, out_shape=(ModelConstants.NUM_LANE_LINES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH))
|
||||
self.parse_mdn('road_edges', outs, out_shape=(ModelConstants.NUM_ROAD_EDGES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH))
|
||||
self.parse_mdn('pose', outs, out_shape=(ModelConstants.POSE_WIDTH,))
|
||||
self.parse_mdn('road_transform', outs, out_shape=(ModelConstants.POSE_WIDTH,))
|
||||
if 'sim_pose' in outs:
|
||||
self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
|
||||
self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,))
|
||||
self.parse_mdn('lead', outs, in_N=lead_in, out_N=lead_out,
|
||||
out_shape=(ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH))
|
||||
self.parse_mdn('sim_pose', outs, out_shape=(ModelConstants.POSE_WIDTH,))
|
||||
self.parse_mdn('wide_from_device_euler', outs, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,))
|
||||
self.parse_mdn('lead', outs, out_shape=(ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH))
|
||||
if 'lat_planner_solution' in outs:
|
||||
self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N,ModelConstants.LAT_PLANNER_SOLUTION_WIDTH))
|
||||
self.parse_mdn('lat_planner_solution', outs, out_shape=(ModelConstants.IDX_N, ModelConstants.LAT_PLANNER_SOLUTION_WIDTH))
|
||||
if 'desired_curvature' in outs:
|
||||
self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,))
|
||||
self.parse_mdn('desired_curvature', outs, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,))
|
||||
for k in ['lead_prob', 'lane_lines_prob', 'meta']:
|
||||
self.parse_binary_crossentropy(k, outs)
|
||||
self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,))
|
||||
self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH))
|
||||
self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN, ModelConstants.DESIRE_PRED_WIDTH))
|
||||
return outs
|
||||
|
||||
@@ -22,16 +22,8 @@ def softmax(x, axis=-1):
|
||||
|
||||
|
||||
class Parser:
|
||||
def __init__(self, ignore_missing=False, mhp_config=None):
|
||||
def __init__(self, ignore_missing=False):
|
||||
self.ignore_missing = ignore_missing
|
||||
# Optional overrides for plan/lead ``in_N`` / ``out_N``. ``None`` or any
|
||||
# missing key falls back to ``SplitModelConstants`` so previously compiled
|
||||
# pkls continue to behave identically.
|
||||
self.mhp = mhp_config or {}
|
||||
|
||||
def _mhp(self, head, default_in, default_out):
|
||||
return (self.mhp.get(f'{head}_mhp_n', default_in),
|
||||
self.mhp.get(f'{head}_mhp_selection', default_out))
|
||||
|
||||
def check_missing(self, outs, name):
|
||||
if name not in outs and not self.ignore_missing:
|
||||
@@ -63,46 +55,36 @@ class Parser:
|
||||
pred_std = safe_exp(raw[:,:,n_values: 2*n_values])
|
||||
|
||||
if in_N > 1:
|
||||
if out_N > 0:
|
||||
weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype)
|
||||
for i in range(out_N):
|
||||
weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1)
|
||||
weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype)
|
||||
for i in range(out_N):
|
||||
weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1)
|
||||
|
||||
if out_N == 1:
|
||||
for fidx in range(weights.shape[0]):
|
||||
idxs = np.argsort(weights[fidx][:,0])[::-1]
|
||||
weights[fidx] = weights[fidx][idxs]
|
||||
pred_mu[fidx] = pred_mu[fidx][idxs]
|
||||
pred_std[fidx] = pred_std[fidx][idxs]
|
||||
assert out_shape is not None
|
||||
full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
|
||||
outs[name + '_weights'] = weights
|
||||
outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
|
||||
outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
|
||||
|
||||
pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
|
||||
pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
|
||||
if out_N == 1:
|
||||
for fidx in range(weights.shape[0]):
|
||||
for hidx in range(out_N):
|
||||
idxs = np.argsort(weights[fidx,:,hidx])[::-1]
|
||||
pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]]
|
||||
pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]]
|
||||
else:
|
||||
# MHP without weights: keep every hypothesis intact.
|
||||
assert out_shape is not None
|
||||
full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
|
||||
outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
|
||||
outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
|
||||
pred_mu_final = pred_mu
|
||||
pred_std_final = pred_std
|
||||
idxs = np.argsort(weights[fidx][:,0])[::-1]
|
||||
weights[fidx] = weights[fidx][idxs]
|
||||
pred_mu[fidx] = pred_mu[fidx][idxs]
|
||||
pred_std[fidx] = pred_std[fidx][idxs]
|
||||
assert out_shape is not None
|
||||
full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
|
||||
outs[name + '_weights'] = weights
|
||||
outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
|
||||
outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
|
||||
|
||||
pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
|
||||
pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
|
||||
for fidx in range(weights.shape[0]):
|
||||
for hidx in range(out_N):
|
||||
idxs = np.argsort(weights[fidx,:,hidx])[::-1]
|
||||
pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]]
|
||||
pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]]
|
||||
else:
|
||||
pred_mu_final = pred_mu
|
||||
pred_std_final = pred_std
|
||||
|
||||
if out_N > 1 or (in_N > 1 and out_N == 0):
|
||||
if out_N > 1:
|
||||
assert out_shape is not None
|
||||
n_selections = out_N if out_N > 1 else in_N
|
||||
final_shape = tuple([raw.shape[0], n_selections] + list(out_shape))
|
||||
final_shape = tuple([raw.shape[0], out_N] + list(out_shape))
|
||||
else:
|
||||
assert out_shape is not None
|
||||
final_shape = tuple([raw.shape[0],] + list(out_shape))
|
||||
@@ -118,26 +100,15 @@ class Parser:
|
||||
|
||||
def parse_dynamic_outputs(self, outs: dict[str, np.ndarray]) -> None:
|
||||
if 'lead' in outs:
|
||||
# Prefer explicit overrides in ``mhp``; otherwise fall back to the
|
||||
# legacy `is_mhp` heuristic that inspects the raw tensor's last axis.
|
||||
if self.mhp.get('lead_mhp_n') is not None or self.mhp.get('lead_mhp_selection') is not None:
|
||||
lead_in_N, lead_out_N = (SplitModelConstants.LEAD_MHP_N, SplitModelConstants.LEAD_MHP_SELECTION)
|
||||
lead_in_N = self.mhp.get('lead_mhp_n', lead_in_N)
|
||||
lead_out_N = self.mhp.get('lead_mhp_selection', lead_out_N)
|
||||
lead_out_shape = (SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH)
|
||||
else:
|
||||
lead_mhp = self.is_mhp(outs, 'lead',
|
||||
SplitModelConstants.LEAD_MHP_SELECTION * SplitModelConstants.LEAD_TRAJ_LEN * SplitModelConstants.LEAD_WIDTH)
|
||||
lead_in_N, lead_out_N = (SplitModelConstants.LEAD_MHP_N, SplitModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0)
|
||||
lead_out_shape = (SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) if lead_mhp else \
|
||||
(SplitModelConstants.LEAD_MHP_SELECTION, SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH)
|
||||
lead_mhp = self.is_mhp(outs, 'lead',
|
||||
SplitModelConstants.LEAD_MHP_SELECTION * SplitModelConstants.LEAD_TRAJ_LEN * SplitModelConstants.LEAD_WIDTH)
|
||||
lead_in_N, lead_out_N = (SplitModelConstants.LEAD_MHP_N, SplitModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0)
|
||||
lead_out_shape = (SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) if lead_mhp else \
|
||||
(SplitModelConstants.LEAD_MHP_SELECTION, SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH)
|
||||
self.parse_mdn('lead', outs, in_N=lead_in_N, out_N=lead_out_N, out_shape=lead_out_shape)
|
||||
if 'plan' in outs:
|
||||
if self.mhp.get('plan_mhp_n') is not None or self.mhp.get('plan_mhp_selection') is not None:
|
||||
plan_in_N, plan_out_N = self._mhp('plan', SplitModelConstants.PLAN_MHP_N, SplitModelConstants.PLAN_MHP_SELECTION)
|
||||
else:
|
||||
plan_mhp = self.is_mhp(outs, 'plan', SplitModelConstants.IDX_N * SplitModelConstants.PLAN_WIDTH)
|
||||
plan_in_N, plan_out_N = (SplitModelConstants.PLAN_MHP_N, SplitModelConstants.PLAN_MHP_SELECTION) if plan_mhp else (0, 0)
|
||||
plan_mhp = self.is_mhp(outs, 'plan', SplitModelConstants.IDX_N * SplitModelConstants.PLAN_WIDTH)
|
||||
plan_in_N, plan_out_N = (SplitModelConstants.PLAN_MHP_N, SplitModelConstants.PLAN_MHP_SELECTION) if plan_mhp else (0, 0)
|
||||
self.parse_mdn('plan', outs, in_N=plan_in_N, out_N=plan_out_N,
|
||||
out_shape=(SplitModelConstants.IDX_N, SplitModelConstants.PLAN_WIDTH))
|
||||
if 'planplus' in outs:
|
||||
@@ -165,7 +136,7 @@ class Parser:
|
||||
self.parse_mdn('road_edges', outs, in_N=0, out_N=0,
|
||||
out_shape=(SplitModelConstants.NUM_ROAD_EDGES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH))
|
||||
if 'sim_pose' in outs:
|
||||
self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,))
|
||||
self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,))
|
||||
if 'action' in outs:
|
||||
self.parse_mdn('action', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.ACTION_WIDTH,))
|
||||
|
||||
|
||||
@@ -1,171 +1,113 @@
|
||||
"""Tests for the dynamic MDN inference in ``mhp_inference`` and the parser
|
||||
changes that read from ``output_slices`` instead of hardcoded constants.
|
||||
"""
|
||||
"""Tests for the dynamic MDN-hypothesis inference in ``parse_model_outputs``."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from openpilot.sunnypilot.modeld_v2.mhp_inference import (
|
||||
infer_mhp,
|
||||
slice_size,
|
||||
infer_mhp_for_outputs,
|
||||
)
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser
|
||||
|
||||
from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser, _infer_mhp
|
||||
from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
|
||||
|
||||
|
||||
# -- infer_mhp --------------------------------------------------------------
|
||||
# -- _infer_mhp -------------------------------------------------------------
|
||||
|
||||
class TestInferMhp:
|
||||
N_PLAN = ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH # 495
|
||||
N_LEAD = ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH # 24
|
||||
P_PLAN = ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH # 495
|
||||
P_LEAD = ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH # 24
|
||||
|
||||
def test_legacy_plan_preserved(self):
|
||||
# Legacy: 5 hypotheses x (2*495 + 1) = 4955
|
||||
assert infer_mhp(4955, self.N_PLAN, legacy_in_n=5, legacy_out_n=1) == (5, 1)
|
||||
def test_legacy_plan(self):
|
||||
# 5 hypotheses * (2*495 + 1) = 4955
|
||||
assert _infer_mhp(4955, self.P_PLAN) == (5, 1)
|
||||
|
||||
def test_legacy_lead_preserved(self):
|
||||
# Legacy: 2 hypotheses x (2*24 + 3) = 102
|
||||
assert infer_mhp(102, self.N_LEAD, legacy_in_n=2, legacy_out_n=3) == (2, 3)
|
||||
def test_legacy_lead(self):
|
||||
# 2 hypotheses * (2*24 + 3) = 102
|
||||
assert _infer_mhp(102, self.P_LEAD) == (2, 3)
|
||||
|
||||
def test_new_supercombo_plan_single_hypothesis_no_weights(self):
|
||||
# New combined supercombo: 1 hypothesis x (2*495 + 0) = 990
|
||||
assert infer_mhp(990, self.N_PLAN, legacy_in_n=5, legacy_out_n=1) == (1, 0)
|
||||
def test_new_plan(self):
|
||||
# 1 hypothesis * (2*495 + 0) = 990
|
||||
assert _infer_mhp(990, self.P_PLAN) == (1, 0)
|
||||
|
||||
def test_new_supercombo_lead_three_hypotheses_no_weights(self):
|
||||
# New combined supercombo: 3 hypotheses x (2*24 + 0) = 144
|
||||
assert infer_mhp(144, self.N_LEAD, legacy_in_n=2, legacy_out_n=3) == (3, 0)
|
||||
def test_new_lead(self):
|
||||
# 3 hypotheses * (2*24 + 0) = 144
|
||||
assert _infer_mhp(144, self.P_LEAD) == (3, 0)
|
||||
|
||||
def test_out_n_one(self):
|
||||
# 4 hypotheses x (2*495 + 1) = 3964
|
||||
assert infer_mhp(3964, self.N_PLAN, legacy_in_n=5, legacy_out_n=1) == (4, 1)
|
||||
def test_no_mhp_output(self):
|
||||
# lane_lines (528 = 4*132): 1 hypothesis * (2*528 + 0) = 1056; 2 hypotheses * 528+0 = 1056 too — same.
|
||||
# In practice a single-hypothesis MDN with ``out_N == 0`` yields 1*2*prod = 528 * 2 = ... no, lane_lines is
|
||||
# ``in_N=0`` because it's laid out as 2*prod directly (1056 = 4 * 132 * 2 simply). The auto-detect prefers
|
||||
# the smallest valid ``in_N``, so 1056 = 2*1*528 + 0 -> (1, 0).
|
||||
assert _infer_mhp(1056, 528) == (1, 0)
|
||||
|
||||
def test_out_n_three(self):
|
||||
# 2 hypotheses x (2*24 + 3) = 102 (matches legacy_lead as well)
|
||||
assert infer_mhp(102, self.N_LEAD, legacy_in_n=4, legacy_out_n=99) == (2, 3)
|
||||
|
||||
def test_zero_or_invalid_returns_fallback(self):
|
||||
assert infer_mhp(0, self.N_PLAN, 5, 1) == (1, 0)
|
||||
assert infer_mhp(-1, self.N_PLAN, 5, 1) == (1, 0)
|
||||
assert infer_mhp(990, 0, 5, 1) == (1, 0)
|
||||
|
||||
def test_no_match_returns_single_hypothesis(self):
|
||||
# 987 doesn't cleanly factor under the constraints we care about.
|
||||
assert infer_mhp(987, self.N_PLAN, 5, 1) == (1, 0)
|
||||
|
||||
|
||||
class TestSliceSize:
|
||||
def test_none_returns_zero(self):
|
||||
assert slice_size(None) == 0
|
||||
|
||||
def test_basic_slice(self):
|
||||
assert slice_size(slice(10, 50)) == 40
|
||||
|
||||
def test_negative_stop_returns_zero(self):
|
||||
assert slice_size(slice(10, -2)) == 0
|
||||
|
||||
def test_none_bounds(self):
|
||||
assert slice_size(slice(None, 100)) == 100
|
||||
|
||||
|
||||
class TestInferMhpForOutputs:
|
||||
def test_infers_for_plan_and_lead(self):
|
||||
slices = {
|
||||
'plan': slice(1576, 2566), # 990
|
||||
'lead': slice(917, 1061), # 144
|
||||
}
|
||||
cfg = infer_mhp_for_outputs(slices, ModelConstants)
|
||||
assert cfg == {'plan_mhp_n': 1, 'plan_mhp_selection': 0,
|
||||
'lead_mhp_n': 3, 'lead_mhp_selection': 0}
|
||||
|
||||
def test_legacy_falls_back_to_constants(self):
|
||||
# Legacy sizes: 4955 plan, 102 lead -> both match Priority 1.
|
||||
slices = {
|
||||
'plan': slice(0, 4955),
|
||||
'lead': slice(4955, 5057),
|
||||
}
|
||||
cfg = infer_mhp_for_outputs(slices, ModelConstants)
|
||||
assert cfg == {'plan_mhp_n': ModelConstants.PLAN_MHP_N,
|
||||
'plan_mhp_selection': ModelConstants.PLAN_MHP_SELECTION,
|
||||
'lead_mhp_n': ModelConstants.LEAD_MHP_N,
|
||||
'lead_mhp_selection': ModelConstants.LEAD_MHP_SELECTION}
|
||||
|
||||
def test_missing_outputs_are_skipped(self):
|
||||
cfg = infer_mhp_for_outputs({}, ModelConstants)
|
||||
assert cfg == {}
|
||||
def test_unknown_size_keeps_single_hypothesis(self):
|
||||
# 989 doesn't divide cleanly under any out_N ∈ {0..6} for P = 495, so we
|
||||
# fall back to the safe single-hypothesis default.
|
||||
assert _infer_mhp(989, self.P_PLAN) == (1, 0)
|
||||
|
||||
|
||||
# -- CombinedParser ---------------------------------------------------------
|
||||
|
||||
def _synth_outputs(in_n_plan=5, out_n_plan=1, in_n_lead=2, out_n_lead=3,
|
||||
n_plan=ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH,
|
||||
n_lead=ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH):
|
||||
"""Synthesize a flat-output dict with the right per-head sizes."""
|
||||
plan_size = in_n_plan * (2 * n_plan + out_n_plan)
|
||||
lead_size = in_n_lead * (2 * n_lead + out_n_lead)
|
||||
return {
|
||||
'plan': np.random.RandomState(0).randn(1, plan_size).astype(np.float32),
|
||||
'lead': np.random.RandomState(1).randn(1, lead_size).astype(np.float32),
|
||||
# Other outputs that parse_outputs() consumes:
|
||||
'lane_lines': np.random.RandomState(2).randn(1, 528).astype(np.float32),
|
||||
'road_edges': np.random.RandomState(3).randn(1, 264).astype(np.float32),
|
||||
'pose': np.random.RandomState(4).randn(1, 12).astype(np.float32),
|
||||
'road_transform': np.random.RandomState(5).randn(1, 12).astype(np.float32),
|
||||
'wide_from_device_euler': np.random.RandomState(6).randn(1, 6).astype(np.float32),
|
||||
'lead_prob': np.random.RandomState(7).randn(1, 3).astype(np.float32),
|
||||
'lane_lines_prob':np.random.RandomState(8).randn(1, 8).astype(np.float32),
|
||||
'meta': np.random.RandomState(9).randn(1, 55).astype(np.float32),
|
||||
'desire_state': np.random.RandomState(10).randn(1, 8).astype(np.float32),
|
||||
'desire_pred': np.random.RandomState(11).randn(1, 32).astype(np.float32),
|
||||
def _synth_outs(
|
||||
plan_in_n: int = 5, plan_out_n: int = 1,
|
||||
lead_in_n: int = 2, lead_out_n: int = 3,
|
||||
extras: bool = True,
|
||||
) -> dict[str, np.ndarray]:
|
||||
plan_size = plan_in_n * (2 * ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH + plan_out_n)
|
||||
lead_size = lead_in_n * (2 * ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH + lead_out_n)
|
||||
rng = np.random.RandomState(0)
|
||||
d = {
|
||||
'plan': rng.randn(1, plan_size).astype(np.float32),
|
||||
'lead': rng.randn(1, lead_size).astype(np.float32),
|
||||
}
|
||||
if extras:
|
||||
d.update({
|
||||
'lane_lines': rng.randn(1, 528).astype(np.float32),
|
||||
'road_edges': rng.randn(1, 264).astype(np.float32),
|
||||
'pose': rng.randn(1, 12).astype(np.float32),
|
||||
'road_transform': rng.randn(1, 12).astype(np.float32),
|
||||
'wide_from_device_euler': rng.randn(1, 6).astype(np.float32),
|
||||
'lead_prob': rng.randn(1, 3).astype(np.float32),
|
||||
'lane_lines_prob':rng.randn(1, 8).astype(np.float32),
|
||||
'meta': rng.randn(1, 55).astype(np.float32),
|
||||
'desire_state': rng.randn(1, 8).astype(np.float32),
|
||||
'desire_pred': rng.randn(1, 32).astype(np.float32),
|
||||
})
|
||||
return d
|
||||
|
||||
|
||||
class TestCombinedParser:
|
||||
def test_legacy_keeps_existing_shape(self):
|
||||
p = CombinedParser() # empty mhp -> legacy constants
|
||||
out = p.parse_outputs(_synth_outputs(5, 1, 2, 3))
|
||||
assert out['plan'].shape == (1, 33, 15)
|
||||
assert out['plan_stds'].shape == (1, 33, 15)
|
||||
# Lead primary output collapses to LEAD_MHP_SELECTION=3 selections per
|
||||
# ``parse_mdn``; raw hypotheses survive as ``lead_hypotheses``.
|
||||
assert out['lead'].shape == (1, 3, 6, 4)
|
||||
assert out['lead_stds'].shape == (1, 3, 6, 4)
|
||||
assert out['plan_hypotheses'].shape == (1, 5, 33, 15)
|
||||
assert out['lead_hypotheses'].shape == (1, 2, 6, 4)
|
||||
def test_legacy_supercombo_shapes(self):
|
||||
p = Parser()
|
||||
out = p.parse_outputs(_synth_outs(plan_in_n=5, plan_out_n=1, lead_in_n=2, lead_out_n=3))
|
||||
assert out['plan'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)
|
||||
assert out['plan_stds'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)
|
||||
assert out['lead'].shape == (1, ModelConstants.LEAD_MHP_SELECTION, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)
|
||||
assert out['lead_stds'].shape == (1, ModelConstants.LEAD_MHP_SELECTION, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)
|
||||
# Per-hypothesis outputs preserved by legacy code path
|
||||
assert out['plan_hypotheses'].shape == (1, 5, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)
|
||||
assert out['lead_hypotheses'].shape == (1, 2, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)
|
||||
|
||||
def test_new_supercombo_plan_and_lead_parse(self):
|
||||
p = CombinedParser(mhp_config={
|
||||
'plan_mhp_n': 1, 'plan_mhp_selection': 0,
|
||||
'lead_mhp_n': 3, 'lead_mhp_selection': 0,
|
||||
})
|
||||
out = p.parse_outputs(_synth_outputs(1, 0, 3, 0))
|
||||
assert out['plan'].shape == (1, 33, 15)
|
||||
assert out['plan_stds'].shape == (1, 33, 15)
|
||||
assert out['lead'].shape == (1, 3, 6, 4)
|
||||
assert out['lead_stds'].shape == (1, 3, 6, 4)
|
||||
# MHP-without-weights keeps every hypothesis as ``*_hypotheses``
|
||||
assert out['lead_hypotheses'].shape == (1, 3, 6, 4)
|
||||
# Plan with a single hypothesis takes the in_N<=1 branch, which
|
||||
# (matching legacy behavior) does not emit ``plan_hypotheses``.
|
||||
assert 'plan_hypotheses' not in out
|
||||
def test_new_supercombo_shapes(self):
|
||||
p = Parser()
|
||||
out = p.parse_outputs(_synth_outs(plan_in_n=1, plan_out_n=0, lead_in_n=3, lead_out_n=0))
|
||||
# Plan single hypothesis collapses straight to (1, IDX_N, PLAN_WIDTH)
|
||||
assert out['plan'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)
|
||||
assert out['plan_stds'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)
|
||||
# Lead with 3 hypotheses and no weights keeps all hypotheses
|
||||
assert out['lead'].shape == (1, 3, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)
|
||||
assert out['lead_stds'].shape == (1, 3, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)
|
||||
assert out['lead_hypotheses'].shape == (1, 3, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)
|
||||
|
||||
|
||||
class TestSplitParser:
|
||||
def test_default_uses_is_mhp_heuristic(self):
|
||||
# No mhp_config -> falls back to inspecting the raw tensor's last axis.
|
||||
# The legacy 102-element lead pack fits the "MHP branch" path.
|
||||
n = SplitParser()
|
||||
outs = {'lead': np.zeros((1, 102), dtype=np.float32)}
|
||||
n.parse_dynamic_outputs(outs)
|
||||
assert outs['lead'].shape == (1, 3, 6, 4)
|
||||
assert outs['lead_hypotheses'].shape == (1, 2, 6, 4)
|
||||
|
||||
def test_explicit_mhp_overrides_is_mhp(self):
|
||||
n = SplitParser(mhp_config={'lead_mhp_n': 3, 'lead_mhp_selection': 0})
|
||||
outs = {'lead': np.zeros((1, 144), dtype=np.float32)}
|
||||
n.parse_dynamic_outputs(outs)
|
||||
assert outs['lead'].shape == (1, 3, 6, 4)
|
||||
def test_unknown_size_does_not_crash_legacy_layout(self):
|
||||
# Provide EVERY output parse_outputs expects so the parser can auto-detect
|
||||
# each head from the actual raw slice size, including the new-supercombo
|
||||
# mixture (legacy-shaped lane_lines/etc. + new-shaped plan/lead).
|
||||
p = Parser()
|
||||
out = p.parse_outputs(_synth_outs(plan_in_n=1, plan_out_n=0, lead_in_n=3, lead_out_n=0))
|
||||
assert out['plan'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)
|
||||
assert out['lead'].shape == (1, 3, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)
|
||||
# Non-MHP outputs must keep the historical 3D shape (no spurious leading 1).
|
||||
assert out['lane_lines'].shape == (1, ModelConstants.NUM_LANE_LINES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH)
|
||||
assert out['pose'].shape == (1, ModelConstants.POSE_WIDTH)
|
||||
assert out['road_transform'].shape == (1, ModelConstants.POSE_WIDTH)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
Reference in New Issue
Block a user