diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index 469110a85..c11f899b7 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -148,7 +148,7 @@ def make_warp_dm(cam_w, cam_h, dm_w, dm_h): stride_pad = stride - cam_w def warp_dm(input_frame, M_inv): - M_inv = M_inv.to(Device.DEFAULT) + M_inv = M_inv.to(Device.DEFAULT).realize() result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w) return result return warp_dm @@ -162,8 +162,14 @@ def make_run_policy(vision_runner, policy_runner, cam_w, cam_h, sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame): - img = shift_and_sample(img_q, frame_prepare(frame, tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn) - big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn) + tfm = tfm.to(Device.DEFAULT) + big_tfm = big_tfm.to(Device.DEFAULT) + desire = desire.to(Device.DEFAULT) + traffic_convention = traffic_convention.to(Device.DEFAULT) + Tensor.realize(tfm, big_tfm, desire, traffic_convention) + + img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn) + big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn) if prepare_only: return img, big_img @@ -172,9 +178,9 @@ def make_run_policy(vision_runner, policy_runner, cam_w, cam_h, new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0) feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn) - desire_buf = shift_and_sample(desire_q, desire.to(Device.DEFAULT).reshape(1, 1, -1), sample_desire_fn) + desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn) - inputs = {'features_buffer': feat_buf, 'desire_pulse': desire_buf, 'traffic_convention': traffic_convention.to(Device.DEFAULT)} + inputs = {'features_buffer': feat_buf, 'desire_pulse': desire_buf, 'traffic_convention': traffic_convention} policy_out = next(iter(policy_runner(inputs).values())).cast('float32') return vision_out, policy_out @@ -210,10 +216,11 @@ def compile_modeld(cam_w, cam_h, prepare_only, pkl_path): def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True): input_queues, npy = make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip) np.random.seed(seed) + Tensor.manual_seed(seed) for i in range(N_RUNS): - frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize() - big_frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize() + frame = Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize() + big_frame = Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize() for v in npy.values(): v[:] = np.random.randn(*v.shape).astype(v.dtype) Device.default.synchronize() @@ -263,7 +270,7 @@ def compile_dm_warp(cam_w, cam_h, pkl_path): warp_dm_jit = TinyJit(warp_dm, prune=True) for i in range(10): - inputs = [Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), + inputs = [Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize(), Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] Device.default.synchronize() st = time.perf_counter()