mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53: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.
157 lines
6.3 KiB
Python
157 lines
6.3 KiB
Python
# much taken from https://github.com/cloneofsimo/minRF
|
|
from tinygrad import Tensor, nn, GlobalCounters, TinyJit, Context
|
|
from tinygrad.helpers import getenv, trange
|
|
from extra.models.llama import Attention, FeedForward, precompute_freqs_cis
|
|
|
|
def modulate(x:Tensor, shift:Tensor, scale:Tensor) -> Tensor: return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
|
|
|
# TODO: why doesn't the TimestepEmbedder from minRF work?
|
|
class TimestepEmbedder:
|
|
def __init__(self, hidden_size): self.mlp = [nn.Linear(1, hidden_size), Tensor.silu, nn.Linear(hidden_size, hidden_size)]
|
|
def __call__(self, t:Tensor): return t.reshape(-1, 1).sequential(self.mlp)
|
|
|
|
class TransformerBlock:
|
|
def __init__(self, dim, n_heads, norm_eps=1e-5):
|
|
self.attention = Attention(dim, n_heads)
|
|
self.feed_forward = FeedForward(dim, 4*dim)
|
|
self.attention_norm = nn.LayerNorm(dim, eps=norm_eps)
|
|
self.ffn_norm = nn.LayerNorm(dim, eps=norm_eps)
|
|
self.adaLN_modulation = nn.Linear(dim, 6 * dim, bias=True)
|
|
|
|
def __call__(self, x:Tensor, freqs_cis:Tensor, adaln_input:Tensor):
|
|
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input.silu()).chunk(6, dim=1)
|
|
x = x + gate_msa.unsqueeze(1) * self.attention(modulate(self.attention_norm(x), shift_msa, scale_msa), 0, freqs_cis)
|
|
x = x + gate_mlp.unsqueeze(1) * self.feed_forward(modulate(self.ffn_norm(x), shift_mlp, scale_mlp))
|
|
return x.contiguous().contiguous_backward()
|
|
|
|
class FinalLayer:
|
|
def __init__(self, dim, patch_size, out_channels):
|
|
self.norm_final = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
|
self.linear = nn.Linear(dim, patch_size*patch_size*out_channels, bias=True)
|
|
self.adaLN_modulation = nn.Linear(dim, 2 * dim, bias=True)
|
|
|
|
# init weights/bias to 0
|
|
self.linear.weight.replace(self.linear.weight.zeros_like().contiguous())
|
|
self.linear.bias.replace(self.linear.bias.zeros_like().contiguous())
|
|
|
|
def __call__(self, x:Tensor, c:Tensor):
|
|
shift, scale = self.adaLN_modulation(c.silu()).chunk(2, dim=1)
|
|
x = modulate(self.norm_final(x), shift, scale)
|
|
return self.linear(x)
|
|
|
|
# channels=1, input_size=32, dim=64, n_layers=6, n_heads=4, num_classes=10
|
|
class DiT_Llama:
|
|
def __init__(self, in_channels=1, dim=64, n_layers=6, n_heads=4, num_classes=10, patch_size=2):
|
|
self.patch_size = patch_size
|
|
self.out_channels = in_channels
|
|
self.num_classes = num_classes
|
|
|
|
self.init_conv_seq = [
|
|
nn.Conv2d(in_channels, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
|
|
nn.Conv2d(dim //2, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
|
|
]
|
|
|
|
self.x_embedder = nn.Linear(self.patch_size * self.patch_size * dim // 2, dim, bias=True)
|
|
self.t_embedder = TimestepEmbedder(dim)
|
|
self.y_embedder = nn.Embedding(num_classes+1, dim)
|
|
self.final_layer = FinalLayer(dim, self.patch_size, self.out_channels)
|
|
|
|
self.freqs_cis = precompute_freqs_cis(dim // n_heads, 4096)
|
|
self.layers = [TransformerBlock(dim, n_heads) for _ in range(n_layers)]
|
|
|
|
def unpatchify(self, x:Tensor):
|
|
c, p = self.out_channels, self.patch_size
|
|
h = w = int(x.shape[1] ** 0.5)
|
|
x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
|
|
x = x.rearrange("n h w p q c -> n c h p w q")
|
|
return x.reshape(shape=(x.shape[0], c, h * p, h * p))
|
|
|
|
def patchify(self, x:Tensor):
|
|
B, C, H, W = x.shape
|
|
x = x.reshape(B, C, H // self.patch_size, self.patch_size, W // self.patch_size, self.patch_size)
|
|
x = x.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
|
|
return x # B <H*W ish> <C*patch_size*patch_size>
|
|
|
|
def __call__(self, x:Tensor, t:Tensor, y:Tensor) -> Tensor:
|
|
x = x.sequential(self.init_conv_seq)
|
|
x = self.patchify(x)
|
|
x = self.x_embedder(x)
|
|
adaln_input = self.t_embedder(t) + self.y_embedder(y)
|
|
adaln_input = adaln_input.contiguous()
|
|
for layer in self.layers:
|
|
x = layer(x, self.freqs_cis[:, :x.size(1)], adaln_input=adaln_input)
|
|
x = self.final_layer(x, adaln_input)
|
|
return self.unpatchify(x)
|
|
|
|
def rf(self, x:Tensor, cond:Tensor):
|
|
b = x.shape[0]
|
|
# self.ln is True
|
|
t = Tensor.randn((b,)).sigmoid()
|
|
texp = t.view([b, *([1] * len(x.shape[1:]))])
|
|
|
|
# conditional dropout
|
|
dropout_prob = 0.1
|
|
cond = (Tensor.rand(cond.shape[0]) < dropout_prob).where(cond.full_like(self.num_classes), cond)
|
|
|
|
# this is rectified flow
|
|
z1 = x.randn_like()
|
|
zt = (1 - texp) * x + texp * z1
|
|
vtheta = self(zt, t, cond)
|
|
|
|
# MSE loss
|
|
return ((z1 - x) - vtheta).square().mean()
|
|
|
|
def sample(self, z, cond, null_cond, sample_steps=50, cfg=2.0):
|
|
b = z.size(0)
|
|
dt = Tensor.full((b,)+(1,)*len(z.shape[1:]), fill_value=1.0/sample_steps).contiguous()
|
|
images = [z]
|
|
for i in range(sample_steps, 0, -1):
|
|
t = Tensor.full((b,), fill_value=i/sample_steps).contiguous()
|
|
vc = self(z, t, cond)
|
|
vu = self(z, t, null_cond)
|
|
vc = vu + cfg * (vc - vu)
|
|
z = z - dt * vc
|
|
z = z.contiguous()
|
|
images.append(z)
|
|
return images
|
|
|
|
def mviz(t:Tensor):
|
|
assert len(t.shape) == 4 and t.shape[1] == 1
|
|
ft = t.permute(1,2,0,3).reshape(32, -1)
|
|
assert ft.shape[-1]%32 == 0
|
|
print("")
|
|
for y in ((ft+1)/2).clamp(0,1).tolist():
|
|
ln = [f"\033[38;5;{232+int(x*23)}m██" for x in y]
|
|
print(''.join(ln) + "\033[0m")
|
|
|
|
if __name__ == "__main__":
|
|
X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
|
|
X_train = X_train.pad((2,2,2,2))
|
|
X_train = ((X_train.float()/255)-0.5)/0.5
|
|
Y_train = Y_train.int()
|
|
|
|
model = DiT_Llama(patch_size=getenv("PATCH_SIZE", 2))
|
|
for r in nn.state.get_parameters(model): r.realize()
|
|
optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=5e-4)
|
|
|
|
@TinyJit
|
|
@Context(TRAINING=1)
|
|
def train_step():
|
|
if getenv("OVERFIT"): samples = Tensor.zeros(getenv("BS", 256), dtype='int')
|
|
else: samples = Tensor.randint(getenv("BS", 256), high=X_train.shape[0])
|
|
optimizer.zero_grad()
|
|
loss = model.rf(X_train[samples], Y_train[samples])
|
|
loss.backward()
|
|
optimizer.step()
|
|
return loss
|
|
|
|
@TinyJit
|
|
def sample(z:Tensor, cond:Tensor) -> Tensor:
|
|
return model.sample(z, cond, Tensor.full_like(cond, 10), sample_steps=getenv("SAMPLE_STEPS", 20))[-1]
|
|
|
|
for steps in (t:=trange(getenv("STEPS", 5000))):
|
|
if steps%10 == 0: mviz(sample(Tensor.randn(3, 1, 32, 32), Tensor([5,0,4], dtype='int')))
|
|
GlobalCounters.reset()
|
|
loss = train_step()
|
|
t.set_description(f"loss: {loss.item():9.2f}")
|