mawei/openpilot-evo
1
0
Fork 0
mirror of https://github.com/infiniteCable2/openpilot.git synced 2026-08-05 08:16:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

modeld: less cross device copies (#38168)

Browse Source 
* single copy for frames

* single copy out

* missed this

* pack npy inputs to the policy

* strict zip

* strict zip
This commit is contained in:
Armand du Parc Locmaria
2026-06-10 22:02:28 -07:00
committed by GitHub
parent 98e5f547ea
commit b2bb71b4b7
2 changed files with 29 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
selfdrive/modeld/compile_modeld.py
+25 -21
View File
@@ -1,6 +1,7 @@
#!/usr/bin/env python3
import argparse
import atexit
import math
import os
import pickle
import time
@@ -33,7 +34,7 @@ from tinygrad.engine.jit import TinyJit
NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size'])
WARP_INPUTS = ['img_q', 'big_img_q', 'tfm', 'big_tfm']
POLICY_INPUTS = ['feat_q', 'desire_q', 'desire', 'traffic_convention', 'action_t']
POLICY_INPUTS = ['feat_q', 'desire_q', 'packed_npy_inputs']
UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32)
UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX)
@@ -130,25 +131,28 @@ def make_warp_input_queues(vision_input_shapes, frame_skip, device):
return input_queues, npy
def get_policy_npy_shapes(policy_input_shapes):
dp = policy_input_shapes['desire_pulse'] # (1, 25, 8)
tc = policy_input_shapes['traffic_convention'] # (1, 2)
#TODO action_t is hardcoded to match tc for future compatibility
shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(tc)}
return shapes, [math.prod(s) for s in shapes.values()]
def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device):
input_queues, npy = make_warp_input_queues(vision_input_shapes, frame_skip, device)
fb = policy_input_shapes['features_buffer'] # (1, 25, 512)
dp = policy_input_shapes['desire_pulse'] # (1, 25, 8)
tc = policy_input_shapes['traffic_convention'] # (1, 2)
#TODO action_t is hardcoded to match tc for future compatibility
at = tc
policy_npy = {
'desire': np.zeros(dp[2], dtype=np.float32),
'traffic_convention': np.zeros(tc, dtype=np.float32),
'action_t': np.zeros(at, dtype=np.float32),
}
npy.update(policy_npy)
shapes, sizes = get_policy_npy_shapes(policy_input_shapes)
packed_npy_inputs = np.zeros(sum(sizes), dtype=np.float32)
# views into the packed inputs, to be refilled at runtime
npy.update({k: v.reshape(s) for (k, s), v in zip(shapes.items(), np.split(packed_npy_inputs, np.cumsum(sizes[:-1])), strict=True)})
input_queues.update({
'feat_q': Tensor(np.zeros((frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]), dtype=np.float32), device=device).contiguous().realize(),
'desire_q': Tensor(np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32), device=device).contiguous().realize(),
**{k: Tensor(v, device='NPY').realize() for k, v in policy_npy.items()},
'packed_npy_inputs': Tensor(packed_npy_inputs, device='NPY').realize(),
})
return input_queues, npy
@@ -175,10 +179,11 @@ def make_warp(nv12, model_w, model_h, frame_skip):
big_tfm = big_tfm.to(WARP_DEV)
Tensor.realize(tfm, big_tfm)
warped_frame = frame_prepare(frame, tfm).unsqueeze(0).to(Device.DEFAULT)
warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0).to(Device.DEFAULT)
img = shift_and_sample(img_q, warped_frame, sample_skip_fn)
big_img = shift_and_sample(big_img_q, warped_big_frame, sample_skip_fn)
warped_frame = frame_prepare(frame, tfm).unsqueeze(0)
warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0)
warped = Tensor.cat(warped_frame, warped_big_frame).to(Device.DEFAULT)
img = shift_and_sample(img_q, warped[0:1], sample_skip_fn)
big_img = shift_and_sample(big_img_q, warped[1:2], sample_skip_fn)
return img, big_img
return warp_enqueue
@@ -187,12 +192,11 @@ def make_run_policy(model_runners, model_metadata, frame_skip):
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
vision_features_slice = model_metadata['vision']['output_slices']['hidden_state']
npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['on_policy']['input_shapes'])
def run_policy(img, big_img, feat_q, desire_q, desire, traffic_convention, action_t):
desire = desire.to(Device.DEFAULT)
traffic_convention = traffic_convention.to(Device.DEFAULT)
action_t = action_t.to(Device.DEFAULT)
Tensor.realize(desire, traffic_convention, action_t)
def run_policy(img, big_img, feat_q, desire_q, packed_npy_inputs):
packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT).realize()
desire, traffic_convention, action_t = (t.reshape(s) for t, s in zip(packed_npy_inputs.split(npy_sizes), npy_shapes.values(), strict=True))
desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
vision_out = next(iter(model_runners['vision']({'img': img, 'big_img': big_img}).values())).cast('float32')
@@ -207,7 +211,7 @@ def make_run_policy(model_runners, model_metadata, frame_skip):
}
on_policy_out = next(iter(model_runners['on_policy'](inputs).values())).cast('float32')
off_policy_out = next(iter(model_runners['off_policy'](inputs).values())).cast('float32')
return vision_out, on_policy_out, off_policy_out
return Tensor.cat(vision_out, on_policy_out, off_policy_out, dim=1),
return run_policy
selfdrive/modeld/modeld.py
+4 -4
View File
@@ -83,6 +83,7 @@ class ModelState:
self.vision_input_shapes = vision_metadata['input_shapes']
self.vision_input_names = list(self.vision_input_shapes.keys())
self.vision_output_slices = vision_metadata['output_slices']
self.vision_output_len = vision_metadata['output_shapes']['outputs'][1]
off_policy_metadata = jits['metadata']['off_policy']
self.off_policy_output_slices = off_policy_metadata['output_slices']
@@ -90,6 +91,7 @@ class ModelState:
policy_metadata = jits['metadata']['on_policy']
self.policy_input_shapes = policy_metadata['input_shapes']
self.policy_output_slices = policy_metadata['output_slices']
self.on_policy_output_len = policy_metadata['output_shapes']['outputs'][1]
self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
@@ -131,13 +133,11 @@ class ModelState:
if prepare_only:
return None
vision_output, on_policy_output, off_policy_output = self.run_policy(
outs, = self.run_policy(
**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, img=img, big_img=big_img
)
vision_output, on_policy_output, off_policy_output = np.split(outs.numpy()[0], [self.vision_output_len, self.vision_output_len+self.on_policy_output_len])
vision_output = vision_output.numpy().flatten()
off_policy_output = off_policy_output.numpy().flatten()
on_policy_output = on_policy_output.numpy().flatten()
vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices))
off_policy_outputs_dict = self.parser.parse_off_policy_outputs(self.slice_outputs(off_policy_output, self.off_policy_output_slices))
policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(on_policy_output, self.policy_output_slices))