mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-11 15:03: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.
84 lines
2.3 KiB
Python
84 lines
2.3 KiB
Python
import unittest, time
|
|
import numpy as np
|
|
from tinygrad import Device
|
|
from tinygrad.nn.state import get_parameters
|
|
from tinygrad.nn import optim
|
|
from tinygrad.helpers import getenv
|
|
from test.helpers import slow
|
|
from extra.training import train
|
|
from extra.models.convnext import ConvNeXt
|
|
from extra.models.efficientnet import EfficientNet
|
|
from extra.models.transformer import Transformer
|
|
from extra.models.vit import ViT
|
|
from extra.models.resnet import ResNet18
|
|
|
|
BS = getenv("BS", 2)
|
|
|
|
def train_one_step(model,X,Y):
|
|
params = get_parameters(model)
|
|
pcount = 0
|
|
for p in params:
|
|
pcount += np.prod(p.shape)
|
|
optimizer = optim.SGD(params, lr=0.001)
|
|
print("stepping %r with %.1fM params bs %d" % (type(model), pcount/1e6, BS))
|
|
st = time.time()
|
|
train(model, X, Y, optimizer, steps=1, BS=BS)
|
|
et = time.time()-st
|
|
print("done in %.2f ms" % (et*1000.))
|
|
|
|
def check_gc():
|
|
if Device.DEFAULT == "CL":
|
|
from extra.introspection import print_objects
|
|
assert print_objects() == 0
|
|
|
|
class TestTrain(unittest.TestCase):
|
|
def test_convnext(self):
|
|
model = ConvNeXt(depths=[1], dims=[16])
|
|
X = np.zeros((BS,3,224,224), dtype=np.float32)
|
|
Y = np.zeros((BS), dtype=np.int32)
|
|
train_one_step(model,X,Y)
|
|
check_gc()
|
|
|
|
@slow
|
|
def test_efficientnet(self):
|
|
model = EfficientNet(0)
|
|
X = np.zeros((BS,3,224,224), dtype=np.float32)
|
|
Y = np.zeros((BS), dtype=np.int32)
|
|
train_one_step(model,X,Y)
|
|
check_gc()
|
|
|
|
@slow
|
|
def test_vit(self):
|
|
model = ViT()
|
|
X = np.zeros((BS,3,224,224), dtype=np.float32)
|
|
Y = np.zeros((BS,), dtype=np.int32)
|
|
train_one_step(model,X,Y)
|
|
check_gc()
|
|
|
|
@slow
|
|
def test_transformer(self):
|
|
# this should be small GPT-2, but the param count is wrong
|
|
# (real ff_dim is 768*4)
|
|
model = Transformer(syms=10, maxlen=6, layers=12, embed_dim=768, num_heads=12, ff_dim=768//4)
|
|
X = np.zeros((BS,6), dtype=np.float32)
|
|
Y = np.zeros((BS,6), dtype=np.int32)
|
|
train_one_step(model,X,Y)
|
|
check_gc()
|
|
|
|
@slow
|
|
def test_resnet(self):
|
|
X = np.zeros((BS, 3, 224, 224), dtype=np.float32)
|
|
Y = np.zeros((BS), dtype=np.int32)
|
|
for resnet_v in [ResNet18]:
|
|
model = resnet_v()
|
|
model.load_from_pretrained()
|
|
train_one_step(model, X, Y)
|
|
check_gc()
|
|
|
|
def test_bert(self):
|
|
# TODO: write this
|
|
pass
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|