mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 15:33: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.
108 lines
3.9 KiB
Python
108 lines
3.9 KiB
Python
import itertools
|
|
from typing import Callable
|
|
from tinygrad import nn, Tensor, dtypes, Device, TinyJit, Context
|
|
from tinygrad.helpers import getenv, trange, partition
|
|
|
|
class Model:
|
|
def __init__(self):
|
|
self.layers: list[Callable[[Tensor], Tensor]] = [
|
|
nn.Conv2d(1, 32, 5), Tensor.relu,
|
|
nn.Conv2d(32, 32, 5), Tensor.relu,
|
|
nn.BatchNorm(32), Tensor.max_pool2d,
|
|
nn.Conv2d(32, 64, 3), Tensor.relu,
|
|
nn.Conv2d(64, 64, 3), Tensor.relu,
|
|
nn.BatchNorm(64), Tensor.max_pool2d,
|
|
lambda x: x.flatten(1), nn.Linear(576, 10)]
|
|
|
|
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
|
|
|
|
# TODO: refactor this into optim/onnx
|
|
def functional_adam(g:Tensor, m:Tensor, v:Tensor, b1_t:Tensor, b2_t:Tensor, lr=0.001, b1=0.9, b2=0.999, eps=1e-6) -> Tensor:
|
|
b1_t *= b1
|
|
b2_t *= b2
|
|
m.assign(b1 * m + (1.0 - b1) * g)
|
|
v.assign(b2 * v + (1.0 - b2) * (g * g))
|
|
m_hat = m / (1.0 - b1_t)
|
|
v_hat = v / (1.0 - b2_t)
|
|
return lr * (m_hat / (v_hat.sqrt() + eps))
|
|
|
|
if __name__ == "__main__":
|
|
BS = getenv("BS", 512)
|
|
ACC_STEPS = getenv("ACC_STEPS", 8)
|
|
|
|
X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
|
|
model = Model()
|
|
|
|
params = nn.state.get_parameters(model)
|
|
|
|
# init params
|
|
for x in params:
|
|
x.replace(x.contiguous())
|
|
Tensor.realize(*params)
|
|
|
|
# split params (with grads) and buffers (without)
|
|
params, buffers = partition(params, lambda x: x.is_param)
|
|
print(f"params: {len(params)} buffers: {len(buffers)}")
|
|
|
|
# optim params
|
|
pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
|
|
adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
|
|
adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
|
|
adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
|
|
adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
|
|
adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]
|
|
|
|
# create loss and grads. init all state so the JIT works on microbatch
|
|
for x in params: x.assign(x.detach())
|
|
loss = Tensor.zeros(tuple()).contiguous()
|
|
grads = Tensor.zeros(pos_params[-1]).contiguous()
|
|
Tensor.realize(*params, *buffers, *adam_params, loss, grads)
|
|
|
|
@TinyJit
|
|
@Context(TRAINING=1)
|
|
def microbatch():
|
|
samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
|
|
for t in params: t.grad = None
|
|
# divide by ACC_STEPS at the loss
|
|
uloss = (model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]) / ACC_STEPS).backward()
|
|
ugrads = Tensor.cat(*[t.grad.contiguous().flatten() for t in params], dim=0)
|
|
for t in params: t.grad = None
|
|
# concat the grads and assign them
|
|
loss.assign(loss + uloss)
|
|
grads.assign(grads + ugrads)
|
|
Tensor.realize(*params, *buffers, loss, grads)
|
|
|
|
@TinyJit
|
|
def optimizer():
|
|
# run optimizer (on CPU, where adam params live)
|
|
delta = functional_adam(grads.to("CPU"), adam_m, adam_v, adam_b1_t, adam_b2_t)
|
|
|
|
# update the params, copying back the delta one at a time to avoid OOM
|
|
# NOTE: the scheduler is ordering things poorly, all the copies are happening before the adds
|
|
for j,tt in enumerate(params):
|
|
tt.assign(tt.detach() - delta[pos_params[j]:pos_params[j+1]].reshape(tt.shape).to(Device.DEFAULT))
|
|
|
|
# realize everything, zero out loss and grads
|
|
loss.assign(Tensor.zeros_like(loss))
|
|
grads.assign(Tensor.zeros_like(grads))
|
|
Tensor.realize(*params, *adam_params, loss, grads)
|
|
|
|
@TinyJit
|
|
def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
|
|
|
|
test_acc = float('nan')
|
|
for i in (t:=trange(getenv("STEPS", 70))):
|
|
# microbatch sets the gradients
|
|
for _ in range(ACC_STEPS): microbatch()
|
|
|
|
# get the loss before the optimizer clears it
|
|
# this is already realized so this isn't a schedule
|
|
loss_item = loss.item()
|
|
|
|
# run the optimizer
|
|
optimizer()
|
|
|
|
# eval
|
|
if i%10 == 9: test_acc = get_test_acc().item()
|
|
t.set_description(f"loss: {loss_item:6.2f} test_accuracy: {test_acc:5.2f}%")
|