mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-27 16:43:59 +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.
46 lines
1.3 KiB
Python
Executable File
46 lines
1.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from tinygrad.nn.state import get_parameters
|
|
from tinygrad.nn import optim
|
|
from tinygrad.helpers import getenv
|
|
from extra.training import train, evaluate
|
|
from extra.models.resnet import ResNet
|
|
from extra.datasets import fetch_mnist
|
|
|
|
|
|
class ComposeTransforms:
|
|
def __init__(self, trans):
|
|
self.trans = trans
|
|
|
|
def __call__(self, x):
|
|
for t in self.trans:
|
|
x = t(x)
|
|
return x
|
|
|
|
if __name__ == "__main__":
|
|
X_train, Y_train, X_test, Y_test = fetch_mnist()
|
|
X_train = X_train.reshape(-1, 28, 28).astype(np.uint8)
|
|
X_test = X_test.reshape(-1, 28, 28).astype(np.uint8)
|
|
classes = 10
|
|
|
|
TRANSFER = getenv('TRANSFER')
|
|
model = ResNet(getenv('NUM', 18), num_classes=classes)
|
|
if TRANSFER:
|
|
model.load_from_pretrained()
|
|
|
|
lr = 5e-3
|
|
transform = ComposeTransforms([
|
|
lambda x: [Image.fromarray(xx).resize((64, 64)) for xx in x],
|
|
lambda x: np.stack([np.asarray(xx) for xx in x], 0),
|
|
lambda x: x / 255.0,
|
|
lambda x: np.tile(np.expand_dims(x, 1), (1, 3, 1, 1)).astype(np.float32),
|
|
])
|
|
for _ in range(5):
|
|
optimizer = optim.SGD(get_parameters(model), lr=lr, momentum=0.9)
|
|
train(model, X_train, Y_train, optimizer, 100, BS=32, transform=transform)
|
|
evaluate(model, X_test, Y_test, num_classes=classes, transform=transform)
|
|
lr /= 1.2
|
|
print(f'reducing lr to {lr:.7f}')
|