mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-07-16 15:12:27 +08:00
modeld: run the driving model as a single combined onnx (#38173)
This commit is contained in:
Harald Schäfer
+5
-3
@@ -25,7 +25,9 @@ class MetadataOnnxPBParser(OnnxPBParser):
|
||||
def get_checkpoint(f):
|
||||
model = MetadataOnnxPBParser(f).parse()
|
||||
metadata = {prop["key"]: prop["value"] for prop in model["metadata_props"]}
|
||||
return metadata['model_checkpoint'].split('/')[0]
|
||||
# "<uuid>" or ".../<run_uuid>/<step>"; combined models list vision then policy
|
||||
parts = metadata['model_checkpoint'].split('/')
|
||||
return parts[-2] if len(parts) > 1 else parts[0]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -37,8 +39,8 @@ if __name__ == "__main__":
|
||||
master_path = MASTER_PATH + MODEL_PATH + fn
|
||||
if os.path.exists(master_path):
|
||||
master = get_checkpoint(master_path)
|
||||
master_col = f"[{master}](https://reporter.comma.life/experiment/{master})"
|
||||
master_col = f"[{master}](https://reporter.comma.life/{master})"
|
||||
else:
|
||||
master_col = "N/A (new model)"
|
||||
pr = get_checkpoint(BASEDIR + MODEL_PATH + fn)
|
||||
print("|", fn, "|", master_col, "|", f"[{pr}](https://reporter.comma.life/experiment/{pr})", "|")
|
||||
print("|", fn, "|", master_col, "|", f"[{pr}](https://reporter.comma.life/{pr})", "|")
|
||||
|
||||
@@ -86,15 +86,14 @@ frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
|
||||
|
||||
for usbgpu in [False, True] if USBGPU else [False]:
|
||||
target_pkl_path = File(modeld_pkl_path(usbgpu)).abspath
|
||||
file_prefix, cmd_flags = ('big_', usbgpu_tg_flags) if usbgpu else ('', tg_flags)
|
||||
driving_onnx_deps = [p for m in [f'{file_prefix}driving_vision', f'{file_prefix}driving_on_policy']
|
||||
for p in get_existing_chunks(File(f"models/{m}.onnx").abspath)]
|
||||
# BIG_INTO_SMALL=1 builds the default target from the big model, e.g. to test it without a USB GPU
|
||||
file_prefix, cmd_flags = ('big_', usbgpu_tg_flags) if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '', tg_flags)
|
||||
driving_onnx_deps = get_existing_chunks(File(f"models/{file_prefix}driving_supercombo.onnx").abspath)
|
||||
camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
|
||||
cmd = (f'{cmd_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py '
|
||||
f'--model-size {model_w}x{model_h} '
|
||||
f'--camera-resolutions {camera_res_args} '
|
||||
f'--vision-onnx {File(f"models/{file_prefix}driving_vision.onnx").abspath} '
|
||||
f'--on-policy-onnx {File(f"models/{file_prefix}driving_on_policy.onnx").abspath} '
|
||||
f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} '
|
||||
f'--output {target_pkl_path} --frame-skip {frame_skip}')
|
||||
onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
|
||||
chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
|
||||
|
||||
@@ -132,26 +132,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)}
|
||||
def get_policy_npy_shapes(input_shapes):
|
||||
dp = input_shapes['desire_pulse'] # (1, 25, 8)
|
||||
tc = input_shapes['traffic_convention'] # (1, 2)
|
||||
at = input_shapes['action_t'] # (1, 2)
|
||||
fb = input_shapes['features_buffer'] # (1, 24, 512)
|
||||
# TODO prev_feat shouldn't exist and be handled inside the JIT, but corrupt on QCOM for now
|
||||
shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(at), 'prev_feat': (fb[0], fb[2])}
|
||||
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)
|
||||
def make_input_queues(input_shapes, frame_skip, device):
|
||||
input_queues, npy = make_warp_input_queues(input_shapes, frame_skip, device)
|
||||
|
||||
fb = policy_input_shapes['features_buffer'] # (1, 25, 512)
|
||||
dp = policy_input_shapes['desire_pulse'] # (1, 25, 8)
|
||||
fb = input_shapes['features_buffer'] # (1, 24, 512), past features only; the model appends the current frame's feature
|
||||
dp = input_shapes['desire_pulse'] # (1, 25, 8)
|
||||
|
||||
shapes, sizes = get_policy_npy_shapes(policy_input_shapes)
|
||||
shapes, sizes = get_policy_npy_shapes(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(),
|
||||
'feat_q': Tensor(np.zeros((frame_skip * fb[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(),
|
||||
'packed_npy_inputs': Tensor(packed_npy_inputs, device='NPY').realize(),
|
||||
})
|
||||
@@ -189,29 +191,27 @@ def make_warp(nv12, model_w, model_h, frame_skip):
|
||||
return warp_enqueue
|
||||
|
||||
|
||||
def make_run_policy(model_runners, model_metadata, frame_skip):
|
||||
def make_run_policy(model_runner, 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'])
|
||||
npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['input_shapes'])
|
||||
|
||||
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, 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)
|
||||
vision_out = next(iter(model_runners['vision']({'img': img, 'big_img': big_img}).values())).cast('float32')
|
||||
|
||||
new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
|
||||
feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
|
||||
feat_buf = shift_and_sample(feat_q, prev_feat.reshape(1, 1, -1), sample_skip_fn)
|
||||
|
||||
inputs = {
|
||||
'img': img,
|
||||
'big_img': big_img,
|
||||
'features_buffer': feat_buf,
|
||||
'desire_pulse': desire_buf,
|
||||
'traffic_convention': traffic_convention,
|
||||
'action_t': action_t,
|
||||
}
|
||||
on_policy_out = next(iter(model_runners['on_policy'](inputs).values())).cast('float32')
|
||||
return Tensor.cat(vision_out, on_policy_out, dim=1),
|
||||
out = next(iter(model_runner(inputs).values())).cast('float32')
|
||||
return out,
|
||||
return run_policy
|
||||
|
||||
|
||||
@@ -281,26 +281,21 @@ if __name__ == "__main__":
|
||||
p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
|
||||
p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True,
|
||||
help='camera resolutions WxH (one or more)')
|
||||
p.add_argument('--vision-onnx', required=True)
|
||||
p.add_argument('--on-policy-onnx', required=True)
|
||||
p.add_argument('--onnx', required=True)
|
||||
p.add_argument('--output', required=True)
|
||||
p.add_argument('--frame-skip', type=int, required=True)
|
||||
args = p.parse_args()
|
||||
|
||||
model_paths = {
|
||||
'vision': read_file_chunked_to_shm(args.vision_onnx),
|
||||
'on_policy': read_file_chunked_to_shm(args.on_policy_onnx),
|
||||
}
|
||||
model_path = read_file_chunked_to_shm(args.onnx)
|
||||
model_w, model_h = args.model_size
|
||||
|
||||
model_runners = {name: OnnxRunner(path) for name, path in model_paths.items()}
|
||||
out = {'metadata': {name: make_metadata_dict(path) for name, path in model_paths.items()}}
|
||||
model_runner = OnnxRunner(model_path)
|
||||
out = {'metadata': make_metadata_dict(model_path)}
|
||||
|
||||
run_policy_jit = TinyJit(make_run_policy(model_runners, out['metadata'], args.frame_skip), prune=True)
|
||||
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']['vision']['input_shapes'],
|
||||
out['metadata']['on_policy']['input_shapes'], args.frame_skip)
|
||||
make_random_model_inputs = partial(make_random_images, keys=['img', 'big_img'], shape=out['metadata']['vision']['input_shapes']['img'])
|
||||
make_policy_queues = partial(make_input_queues, out['metadata']['input_shapes'], args.frame_skip)
|
||||
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)
|
||||
|
||||
@@ -308,7 +303,7 @@ if __name__ == "__main__":
|
||||
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_enqueue = TinyJit(make_warp(nv12, model_w, model_h, args.frame_skip), prune=True)
|
||||
make_warp_queues = partial(make_warp_input_queues, out['metadata']['vision']['input_shapes'], args.frame_skip)
|
||||
make_warp_queues = partial(make_warp_input_queues, out['metadata']['input_shapes'], args.frame_skip)
|
||||
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:
|
||||
|
||||
+10
-16
@@ -79,20 +79,15 @@ class ModelState:
|
||||
input_devices = get_tg_input_devices(PROCESS_NAME, usbgpu)
|
||||
self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV']
|
||||
jits = pickle.loads(read_file_chunked(modeld_pkl_path(usbgpu)))
|
||||
vision_metadata = jits['metadata']['vision']
|
||||
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]
|
||||
|
||||
policy_metadata = jits['metadata']['on_policy']
|
||||
self.policy_input_shapes = policy_metadata['input_shapes']
|
||||
self.policy_output_slices = policy_metadata['output_slices']
|
||||
metadata = jits['metadata']
|
||||
self.input_shapes = metadata['input_shapes']
|
||||
self.vision_input_names = [k for k in self.input_shapes if 'img' in k]
|
||||
self.output_slices = metadata['output_slices']
|
||||
|
||||
self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
|
||||
|
||||
self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
|
||||
self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip, device=self.QUEUE_DEV)
|
||||
self.input_queues, self.npy = make_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
|
||||
self.full_frames: dict[str, Tensor] = {}
|
||||
self._blob_cache: dict[int, Tensor] = {}
|
||||
self.parser = Parser()
|
||||
@@ -132,14 +127,13 @@ class ModelState:
|
||||
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 = np.split(outs.numpy()[0], [self.vision_output_len])
|
||||
vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices))
|
||||
policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(on_policy_output, self.policy_output_slices))
|
||||
combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict}
|
||||
model_output = outs.numpy()[0]
|
||||
outputs_dict = self.parser.parse_outputs(self.slice_outputs(model_output, self.output_slices))
|
||||
self.npy['prev_feat'][:] = model_output[self.output_slices['hidden_state']]
|
||||
|
||||
if SEND_RAW_PRED:
|
||||
combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), on_policy_output.copy()])
|
||||
return combined_outputs_dict
|
||||
outputs_dict['raw_pred'] = model_output.copy()
|
||||
return outputs_dict
|
||||
|
||||
|
||||
def main(demo=False):
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:565e53c38dcd64c50dd3fe4d5ee1530213aeefd66c3f6b67ea6a72a32612a6bf
|
||||
size 14061419
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:471f372751d5931e939320c211e35b7255f8fa8015125b3b3fd48ef43020257e
|
||||
size 195490097
|
||||
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:1f0cab5033fe9e3bc5e174a2e790fa277f7d9fc44c65822d734064d2f899a9a0
|
||||
size 296203378
|
||||
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:78477124cbf3ffe30fa951ebada8410b43c4242c6054584d656f1d329b067e15
|
||||
size 14060847
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:659727c4d4839adc4992a254409a54259a8756a743f2d567bf5fdc6579f8009b
|
||||
size 60881999
|
||||
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:ee29ee5bce84d1ce23e9ff381280de9b4e4d96d2934cd751740354884e112c66
|
||||
size 46877473
|
||||
+1
-1
Submodule tinygrad_repo updated: 556defa0f7...5039d954f2
Reference in New Issue
Block a user