mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +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.
64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
from tinygrad.tensor import Tensor
|
|
from tinygrad.nn import Conv2d, LayerNorm, LayerNorm2d, Linear
|
|
from tinygrad.helpers import fetch, get_child, Context
|
|
|
|
class Block:
|
|
def __init__(self, dim):
|
|
self.dwconv = Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
|
|
self.norm = LayerNorm(dim, eps=1e-6)
|
|
self.pwconv1 = Linear(dim, 4 * dim)
|
|
self.pwconv2 = Linear(4 * dim, dim)
|
|
self.gamma = Tensor.ones(dim)
|
|
|
|
def __call__(self, x:Tensor):
|
|
return x + x.sequential([
|
|
self.dwconv, lambda x: x.permute(0, 2, 3, 1), self.norm,
|
|
self.pwconv1, Tensor.gelu, self.pwconv2, lambda x: (self.gamma * x).permute(0, 3, 1, 2)
|
|
])
|
|
|
|
class ConvNeXt:
|
|
def __init__(self, in_chans=3, num_classes=1000, depths=[3, 3, 9, 3], dims=[96, 192, 384, 768]):
|
|
self.downsample_layers = [
|
|
[Conv2d(in_chans, dims[0], kernel_size=4, stride=4), LayerNorm2d(dims[0], eps=1e-6)],
|
|
*[[LayerNorm2d(dims[i], eps=1e-6), Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2)] for i in range(len(dims)-1)]
|
|
]
|
|
self.stages = [[Block(dims[i]) for _ in range(depths[i])] for i in range(len(dims))]
|
|
self.norm = LayerNorm(dims[-1])
|
|
self.head = Linear(dims[-1], num_classes)
|
|
|
|
def __call__(self, x:Tensor):
|
|
for downsample, stage in zip(self.downsample_layers, self.stages):
|
|
x = x.sequential(downsample).sequential(stage)
|
|
return x.mean([-2, -1]).sequential([self.norm, self.head])
|
|
|
|
# *** model definition is done ***
|
|
|
|
versions = {
|
|
"tiny": {"depths": [3, 3, 9, 3], "dims": [96, 192, 384, 768]},
|
|
"small": {"depths": [3, 3, 27, 3], "dims": [96, 192, 384, 768]},
|
|
"base": {"depths": [3, 3, 9, 3], "dims": [128, 256, 512, 1024]},
|
|
"large": {"depths": [3, 3, 27, 3], "dims": [192, 384, 768, 1536]},
|
|
"xlarge": {"depths": [3, 3, 27, 3], "dims": [256, 512, 1024, 2048]}
|
|
}
|
|
|
|
def get_model(version, load_weights=False):
|
|
model = ConvNeXt(**versions[version])
|
|
if load_weights:
|
|
from tinygrad.nn.state import torch_load
|
|
weights = torch_load(fetch(f'https://dl.fbaipublicfiles.com/convnext/convnext_{version}_1k_224_ema.pth'))['model']
|
|
for k,v in weights.items():
|
|
mv = get_child(model, k)
|
|
mv.assign(v.reshape(mv.shape).to(mv.device)).realize()
|
|
return model
|
|
|
|
if __name__ == "__main__":
|
|
model = get_model("tiny", True)
|
|
|
|
# load image
|
|
from test.models.test_efficientnet import chicken_img, preprocess, _LABELS
|
|
img = Tensor(preprocess(chicken_img))
|
|
|
|
with Context(TRAINING=0):
|
|
out = model(img).numpy()
|
|
print(_LABELS[out.argmax()])
|