mirror of
https://github.com/infiniteCable2/openpilot.git
synced 2026-07-26 03:42:05 +08:00
Revert "modeld: enqueue in policy" (#38266)
Revert "modeld: enqueue in policy (#38264)"
This reverts commit 8c927d0fc2.
This commit is contained in:
committed by
GitHub
parent
8c927d0fc2
commit
8a3bbcdd4e
@@ -34,8 +34,8 @@ from tinygrad.engine.jit import TinyJit
|
||||
|
||||
|
||||
NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size'])
|
||||
WARP_INPUTS = ['tfm', 'big_tfm']
|
||||
POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs']
|
||||
WARP_INPUTS = ['img_q', 'big_img_q', 'tfm', 'big_tfm']
|
||||
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)
|
||||
@@ -175,17 +175,20 @@ def sample_desire(buf, frame_skip):
|
||||
|
||||
def make_warp(nv12, model_w, model_h, frame_skip):
|
||||
frame_prepare = make_frame_prepare(nv12, model_w, model_h)
|
||||
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
||||
|
||||
def warp(tfm, big_tfm, frame, big_frame):
|
||||
def warp_enqueue(img_q, big_img_q, tfm, big_tfm, frame, big_frame):
|
||||
tfm = tfm.to(WARP_DEV)
|
||||
big_tfm = big_tfm.to(WARP_DEV)
|
||||
Tensor.realize(tfm, big_tfm)
|
||||
|
||||
warped_frame = frame_prepare(frame, tfm).unsqueeze(0)
|
||||
warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0)
|
||||
return Tensor.cat(warped_frame, warped_big_frame)
|
||||
|
||||
return warp
|
||||
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
|
||||
|
||||
|
||||
def make_run_policy(model_runner, model_metadata, frame_skip):
|
||||
@@ -193,14 +196,8 @@ def make_run_policy(model_runner, model_metadata, frame_skip):
|
||||
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
||||
npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['input_shapes'])
|
||||
|
||||
def run_policy(warped, img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
|
||||
packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT)
|
||||
warped = warped.to(Device.DEFAULT)
|
||||
Tensor.realize(packed_npy_inputs, warped)
|
||||
|
||||
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)
|
||||
|
||||
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, prev_feat = (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)
|
||||
feat_buf = shift_and_sample(feat_q, prev_feat.reshape(1, 1, -1), sample_skip_fn)
|
||||
@@ -298,16 +295,16 @@ if __name__ == "__main__":
|
||||
run_policy_jit = TinyJit(make_run_policy(model_runner, out['metadata'], args.frame_skip), prune=True)
|
||||
|
||||
make_policy_queues = partial(make_input_queues, out['metadata']['input_shapes'], args.frame_skip)
|
||||
make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, *out['metadata']['input_shapes']['img'][2:]))
|
||||
make_random_model_inputs = partial(make_random_images, keys=['img', 'big_img'], shape=out['metadata']['input_shapes']['img'])
|
||||
out['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS,
|
||||
make_policy_queues)
|
||||
|
||||
for cam_w, cam_h in args.camera_resolutions:
|
||||
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=WARP_DEV)
|
||||
warp = TinyJit(make_warp(nv12, model_w, model_h, args.frame_skip), prune=True)
|
||||
warp_enqueue = TinyJit(make_warp(nv12, model_w, model_h, args.frame_skip), prune=True)
|
||||
make_warp_queues = partial(make_warp_input_queues, out['metadata']['input_shapes'], args.frame_skip)
|
||||
out[(cam_w,cam_h)] = compile_jit(warp, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)
|
||||
out[(cam_w,cam_h)] = compile_jit(warp_enqueue, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)
|
||||
|
||||
with open(args.output, "wb") as f:
|
||||
pickle.dump(out, f)
|
||||
|
||||
@@ -94,7 +94,7 @@ class ModelState:
|
||||
self.parser = Parser()
|
||||
self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')}
|
||||
self.run_policy = jits['run_policy']
|
||||
self.warp = jits[(cam_w,cam_h)]
|
||||
self.warp_enqueue = jits[(cam_w,cam_h)]
|
||||
|
||||
def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
|
||||
parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
|
||||
@@ -120,10 +120,10 @@ class ModelState:
|
||||
self.npy['tfm'][:,:] = transforms['img'][:,:]
|
||||
self.npy['big_tfm'][:,:] = transforms['big_img'][:,:]
|
||||
|
||||
warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'])
|
||||
img, big_img = self.warp_enqueue(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'])
|
||||
|
||||
outs, = self.run_policy(
|
||||
**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped
|
||||
**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, img=img, big_img=big_img
|
||||
)
|
||||
model_output = outs.numpy()[0]
|
||||
outputs_dict = self.parser.parse_outputs(self.slice_outputs(model_output, self.output_slices))
|
||||
|
||||
Reference in New Issue
Block a user