mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-08 22:48:25 -05:00
744af193f0
* remove ExecItem and merge it with ScheduleItem * less diff * fix issues * min diff * don't change bufs in _lower * min diff * update * revert * fixes * diff
57 lines
1.5 KiB
Python
57 lines
1.5 KiB
Python
# abstractions2 goes from back to front, here we will go from front to back
|
|
from typing import List
|
|
from tinygrad.helpers import tqdm
|
|
|
|
# *****
|
|
# 0. Load mnist on the device
|
|
|
|
from tinygrad.nn.datasets import mnist
|
|
X_train, Y_train, _, _ = mnist()
|
|
X_train = X_train.float()
|
|
X_train -= X_train.mean()
|
|
|
|
# *****
|
|
# 1. Define an MNIST model.
|
|
|
|
from tinygrad import Tensor
|
|
|
|
l1 = Tensor.kaiming_uniform(128, 784)
|
|
l2 = Tensor.kaiming_uniform(10, 128)
|
|
def model(x): return x.flatten(1).dot(l1.T).relu().dot(l2.T)
|
|
l1n, l2n = l1.numpy(), l2.numpy()
|
|
|
|
# *****
|
|
# 2. Choose a batch for training and do the backward pass.
|
|
|
|
from tinygrad.nn.optim import SGD
|
|
optim = SGD([l1, l2])
|
|
|
|
Tensor.training = True
|
|
X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
|
|
optim.zero_grad()
|
|
model(X).sparse_categorical_crossentropy(Y).backward()
|
|
optim.schedule_step() # this will step the optimizer without running realize
|
|
|
|
# *****
|
|
# 3. Create a schedule.
|
|
|
|
# The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
|
|
# l1.uop and l2.uop define a computation graph
|
|
|
|
from tinygrad.engine.schedule import ExecItem
|
|
schedule: List[ExecItem] = Tensor.schedule(l1, l2)
|
|
|
|
print(f"The schedule contains {len(schedule)} items.")
|
|
for si in schedule: print(str(si)[:80])
|
|
|
|
# *****
|
|
# 4. Lower and run the schedule.
|
|
|
|
for si in tqdm(schedule): si.run()
|
|
|
|
# *****
|
|
# 5. Print the weight change
|
|
|
|
print("first weight change\n", l1.numpy()-l1n)
|
|
print("second weight change\n", l2.numpy()-l2n)
|