mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
54 lines
2.8 KiB
Python
54 lines
2.8 KiB
Python
import argparse, time
|
|
from contextlib import nullcontext
|
|
import numpy as np
|
|
from tinygrad import Device
|
|
from tinygrad.engine.realize import lower_and_compile
|
|
from extra.bench_log import WallTimeEvent, BenchEvent
|
|
from tinygrad.helpers import getenv
|
|
from tinygrad.nn.state import get_parameters
|
|
from examples.openpilot.helpers import allocate_inputs, load_pickle
|
|
|
|
|
|
def make_inputs(artifact, seed=100):
|
|
rng = np.random.default_rng(seed)
|
|
def initialize(views):
|
|
for value in views.values():
|
|
value[...] = (rng.standard_normal(value.shape) * 8 if np.issubdtype(value.dtype, np.floating) else
|
|
rng.integers(0, 2 if value.dtype == np.bool_ else 16, value.shape))
|
|
inputs = allocate_inputs(artifact['input_specs'], initialize)
|
|
if 'output_specs' in artifact: inputs['output_buffers'] = allocate_inputs(artifact['output_specs'])
|
|
return inputs
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description="Benchmark loading or running compiled model and warp pickles.")
|
|
parser.add_argument('pickle', nargs='?', default='/tmp/openpilot.pkl')
|
|
parser.add_argument('--run', action='store_true', help='benchmark inference instead of loading')
|
|
parser.add_argument('--out-of-band', action='store_true', default=bool(getenv('PICKLE_OOB')))
|
|
parser.add_argument('--runs', type=int, help='defaults to 10 loads or 20 inference runs')
|
|
parser.add_argument('--retarget', action='store_true', help='retarget loaded jit (requires jit to have been compiled with --retargetable)')
|
|
args = parser.parse_args()
|
|
if not args.run:
|
|
load_times = []
|
|
for _ in range(args.runs or 10):
|
|
with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(args.pickle, out_of_band=args.out_of_band)
|
|
load_times.append(wte.time)
|
|
print(f"pickle load: {wte.time:6.2f} s")
|
|
if (limit := getenv("ASSERT_MIN_LOAD_TIME", 0.0)):
|
|
assert min(load_times) < limit, f"Speed regression, expected < {limit} s but took {min(load_times)} s"
|
|
else:
|
|
artifact = load_pickle(args.pickle, out_of_band=args.out_of_band)
|
|
if args.retarget: artifact['run'].captured._linear = lower_and_compile(artifact['run'].captured._linear)
|
|
inputs = make_inputs(artifact)
|
|
times = []
|
|
for _ in range(args.runs or 20):
|
|
start = time.perf_counter()
|
|
with WallTimeEvent(BenchEvent.STEP) if getenv('BENCHMARK_LOG', '') else nullcontext():
|
|
output = artifact['run'](**inputs)
|
|
enqueued = time.perf_counter()
|
|
for device in {tensor.device for tensor in get_parameters((inputs, output))}: Device[device].synchronize()
|
|
times.append((time.perf_counter() - start) * 1e3)
|
|
print(f"enqueue {(enqueued-start)*1e3:6.2f} ms -- total run {times[-1]:6.2f} ms")
|
|
if (limit := getenv("ASSERT_MIN_STEP_TIME", 0.0)):
|
|
assert min(times) < limit, f"Speed regression, expected < {limit} ms but took {min(times)} ms"
|