mawei/sunnypilot
1
0
Fork 0
mirror of https://github.com/sunnypilot/sunnypilot.git synced 2026-08-05 09:05:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

prev_feat in cpu to prevent corruption

Browse Source
This commit is contained in:
discountchubbs
2026-08-04 13:15:12 -07:00
parent 6bfbf45c6d
commit 07bd1dfb52
2 changed files with 72 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
openpilot/sunnypilot/modeld_v2/compile_modeld.py
+65 -30
View File
@@ -57,7 +57,25 @@ def derive_frame_skip(vision_input_shapes: dict, policy_input_shapes: dict) -> i
return 1 if not features_buffer or features_buffer[1] >= 99 else 4
def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]:
def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tuple[dict, list[int]]:
desire_key = _detect_desire_key(input_shapes)
shapes = {}
if desire_key:
shapes['desire'] = (input_shapes[desire_key][2],)
if is_supercombo and 'features_buffer' in input_shapes:
fb = input_shapes['features_buffer']
shapes['prev_feat'] = (fb[0], fb[2])
for key, shape in input_shapes.items():
if key not in (desire_key, 'features_buffer') and 'img' not in key:
shapes[key] = tuple(shape)
sizes = [int(np.prod(size)) for size in shapes.values()]
return shapes, sizes
def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT, is_supercombo: bool = False) -> tuple[dict, dict]:
road_key, _ = _detect_vision_keys(input_shapes)
if not road_key:
raise ValueError("Vision road key missing from input shapes.")
@@ -74,36 +92,41 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D
features_buffer = input_shapes.get('features_buffer')
npy_arrays = {
'desire': np.zeros(desire_shape[2], dtype=np.float32),
'tfm': np.zeros((3, 3), dtype=np.float32),
'big_tfm': np.zeros((3, 3), dtype=np.float32)
}
for key, shape in input_shapes.items():
if key not in npy_arrays and 'img' not in key and key not in ('features_buffer', desire_key):
npy_arrays[key] = np.zeros(shape, dtype=np.float32)
shapes, sizes = get_policy_npy_shapes(input_shapes, is_supercombo=is_supercombo)
packed_npy_inputs = np.zeros(sum(sizes), dtype=np.float32)
split_indices = np.cumsum(sizes[:-1]) if len(sizes) > 1 else []
split_views = np.split(packed_npy_inputs, split_indices) if len(sizes) > 0 else []
for (k, s), v in zip(shapes.items(), split_views, strict=True):
actual_key = desire_key if k == 'desire' else k
npy_arrays[actual_key] = v.reshape(s)
queues = {
'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]),
dtype=np.float32), device=device).contiguous().realize()
dtype=np.float32), device=device).contiguous().realize(),
'packed_npy_inputs': Tensor(packed_npy_inputs, device='NPY').realize(),
}
if features_buffer:
queues['feat_q'] = Tensor(np.zeros((frame_skip * (features_buffer[1] - 1) + 1, features_buffer[0], features_buffer[2]),
dtype=np.float32), device=device).contiguous().realize()
queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items()})
queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items() if key in ('tfm', 'big_tfm')})
return queues, npy_arrays
def make_split_input_queues(vision_input_shapes: dict, policy_input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]:
return generate_queues_and_npy({**vision_input_shapes, **policy_input_shapes}, frame_skip, device)
return generate_queues_and_npy({**vision_input_shapes, **policy_input_shapes}, frame_skip, device, is_supercombo=False)
def make_supercombo_input_queues(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]:
return generate_queues_and_npy(input_shapes, frame_skip, device)
return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True)
def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, model_size: tuple[int, int],
@@ -118,22 +141,17 @@ def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, mode
if not desire_key or not road_key or not wide_key:
raise ValueError("Missing required vision or desire keys in input shapes.")
extra_keys = [key for key in input_shapes if key not in (desire_key, 'features_buffer', 'traffic_convention') and 'img' not in key]
is_supercombo = vision_runner is None
npy_shapes, npy_sizes = get_policy_npy_shapes(input_shapes, is_supercombo=is_supercombo)
def runner(img_q, big_img_q, feat_q, frame, big_frame, tfm, big_tfm, **kwargs):
def runner(img_q, big_img_q, feat_q, packed_npy_inputs, frame, big_frame, tfm, big_tfm, **kwargs):
desire_q = kwargs['desire_q']
desire = kwargs['desire']
traffic_convention = kwargs.get('traffic_convention')
npys = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT), desire.to(Device.DEFAULT)]
if traffic_convention is not None:
npys.append(traffic_convention.to(Device.DEFAULT))
packed_npy_inputs_dev = packed_npy_inputs.to(Device.DEFAULT)
tfm_dev = tfm.to(Device.DEFAULT)
big_tfm_dev = big_tfm.to(Device.DEFAULT)
extra_tensors = {key: kwargs[key].to(Device.DEFAULT) for key in extra_keys if key in kwargs}
Tensor.realize(*npys, *extra_tensors.values())
tfm_dev, big_tfm_dev, desire_dev = npys[:3]
traffic_conv_dev = npys[3] if traffic_convention is not None else None
Tensor.realize(packed_npy_inputs_dev, tfm_dev, big_tfm_dev)
img = shift_and_sample(img_q, frame_prepare(frame, tfm_dev).unsqueeze(0), sample_skip_fn).realize()
big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm_dev).unsqueeze(0), sample_skip_fn).realize()
@@ -141,22 +159,37 @@ def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, mode
if prepare_only:
return img, big_img
desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn).realize()
inputs = {desire_key: desire_buf, **extra_tensors}
unpacked_tensors = [tensor.reshape(shape) for tensor, shape in zip(packed_npy_inputs_dev.split(npy_sizes), npy_shapes.values(), strict=True)]
unpacked_dict = dict(zip(npy_shapes.keys(), unpacked_tensors, strict=True))
if traffic_conv_dev is not None:
inputs['traffic_convention'] = traffic_conv_dev
desire_dev = unpacked_dict['desire']
desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn).realize()
inputs = {desire_key: desire_buf}
for key, tensor_val in unpacked_dict.items():
if key not in ('desire', 'prev_feat'):
inputs[key] = tensor_val
if 'prev_feat' in unpacked_dict:
prev_feat_dev = unpacked_dict['prev_feat']
inputs['features_buffer'] = shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).realize()
if vision_runner:
vision_out_cast = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast('float32').realize()
new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0)
inputs['features_buffer'] = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
if 'features_buffer' not in inputs:
new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0)
inputs['features_buffer'] = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
policy_outs = [next(iter(pol_runner(inputs).values())).cast('float32').realize() for pol_runner in policy_runners]
return (vision_out_cast, *policy_outs) if len(policy_outs) > 1 else (vision_out_cast, policy_outs[0])
inputs.update({road_key: img, wide_key: big_img, 'features_buffer': sample_skip_fn(feat_q)})
inputs.update({road_key: img, wide_key: big_img})
if 'features_buffer' not in inputs:
inputs['features_buffer'] = sample_skip_fn(feat_q)
policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32').realize()
new_feat = policy_out[:, features_slice].reshape(1, -1).unsqueeze(0)
shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
if 'features_buffer' not in inputs and features_slice is not None:
new_feat = policy_out[:, features_slice].reshape(1, -1).unsqueeze(0)
shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
return policy_out
return runner
@@ -192,6 +225,7 @@ def compile_and_warmup(nv12: NV12Frame, model_size: tuple[int, int], prepare_onl
Device.default.synchronize()
print(f" [{i + 1}/3] enqueue {(mid_time - start_time) * 1e3:6.2f} ms -- total {(time.perf_counter() - start_time) * 1e3:6.2f} ms")
# TODO-SP: switch to dump_oob/load_oob on next full recompile of all models
return pickle.loads(pickle.dumps(run_jit)) if not prepare_only else run_jit
@@ -289,6 +323,7 @@ if __name__ == "__main__":
vision_runner, policy_runners, output_data['metadata']))
with open(args.output, "wb") as file:
# TODO-SP: switch to dump_oob from openpilot/selfdrive/helpers on next full recompile of all models
pickle.dump(output_data, file)
pkl_size = os.path.getsize(args.output)
openpilot/sunnypilot/modeld_v2/modeld.py
+7 -1
View File
@@ -108,6 +108,7 @@ class ModelState(ModelStateBase):
def _init_combined(self, pkl_path, cam_w, cam_h, bundle):
cloudlog.warning(f"loading combined pkl: {pkl_path}")
# TODO-SP: switch to load_oob from openpilot/selfdrive/helpers on next full recompile of all models
jits = pickle.load(open_file_chunked(pkl_path))
self.DEV = Device.DEFAULT
@@ -223,11 +224,16 @@ class ModelState(ModelStateBase):
model_output = raw_outputs.numpy().flatten()
sliced = {k: model_output[np.newaxis, v] for k, v in self.vision_output_slices.items()}
outputs = self.parser.parse_outputs(sliced)
if 'prev_feat' in self.numpy_inputs:
self.numpy_inputs['prev_feat'][:] = model_output[self.vision_output_slices['hidden_state']]
else:
vision_output = raw_outputs[0].numpy().flatten()
vision_sliced = {k: vision_output[np.newaxis, v] for k, v in self.vision_output_slices.items()}
outputs = self.parser.parse_vision_outputs(vision_sliced)
if 'prev_feat' in self.numpy_inputs and 'hidden_state' in self.vision_output_slices:
self.numpy_inputs['prev_feat'][:] = vision_output[self.vision_output_slices['hidden_state']]
for i, policy_slices in enumerate(self._policy_slices_list):
policy_output = raw_outputs[i + 1].numpy().flatten()
policy_sliced = {k: policy_output[np.newaxis, v] for k, v in policy_slices.items()}
@@ -248,7 +254,7 @@ class ModelState(ModelStateBase):
buf[0, :-1] = buf[0, 1:]
buf[0, -1, :] = outputs['desired_curvature'][0, :] if not self.mlsim else 0
if 'feat_q' in self.input_queues:
if 'prev_feat' not in self.numpy_inputs and 'feat_q' in self.input_queues:
feat_val = self.input_queues['feat_q'].numpy()
self.input_queues['feat_q'].assign(feat_val).realize()