mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 23:43:44 +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.
50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
import os
|
|
#os.environ["METAL"] = "1"
|
|
import numpy as np
|
|
|
|
BS = 64
|
|
CIN = 256
|
|
COUT = 256
|
|
HW = 32
|
|
K = 3
|
|
PADDING = 0
|
|
# TODO: this is doing some trick, since with CIN=256 COUT=256 it's over 10.4 TFLOPS.
|
|
# are winograd convs less flops? it appears so if they are batched
|
|
# https://www.cse.ust.hk/~weiwa/papers/yan-ppopp20.pdf
|
|
FLOPS = BS*K*K*CIN*HW*HW*COUT*2
|
|
|
|
nb = np.random.default_rng().standard_normal(size=(BS,CIN,HW,HW), dtype=np.float32)
|
|
nc = np.random.default_rng().standard_normal(size=(COUT,CIN,K,K), dtype=np.float32)
|
|
|
|
try:
|
|
import time, torch, torch.mps
|
|
b = torch.from_numpy(nb).to('mps')
|
|
c = torch.from_numpy(nc).to('mps')
|
|
|
|
def torch_prog(b, c):
|
|
st = time.perf_counter()
|
|
a = torch.nn.functional.conv2d(b, c, padding=PADDING)
|
|
torch.mps.synchronize()
|
|
return time.perf_counter() - st
|
|
tm = min([torch_prog(b, c) for _ in range(20)])
|
|
print(f"{tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS conv in torch")
|
|
except RuntimeError:
|
|
print("no torch metal conv")
|
|
|
|
from tinygrad.tensor import Tensor
|
|
from tinygrad.engine.jit import TinyJit
|
|
from tinygrad import Device
|
|
b = Tensor(nb)
|
|
c = Tensor(nc)
|
|
# TODO: slowness without the JIT I suspect comes from a lack of a caching allocator
|
|
@TinyJit
|
|
def tiny_jit(b, c):
|
|
return b.conv2d(c, padding=PADDING).realize()
|
|
def tiny_prog(b, c):
|
|
st = time.perf_counter()
|
|
a = tiny_jit(b, c)
|
|
Device[a.device].synchronize()
|
|
return time.perf_counter() - st
|
|
tm = min([tiny_prog(b, c) for _ in range(5)])
|
|
print(f"{tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS conv in tinygrad")
|