Deep learning фреймворк с нуля

Autograd, свёртки, квантизация, ансамбли. Всё написано руками на NumPy. Не обёртка над PyTorch, а способ понять, как он работает внутри.

Документация Как работает

train.py

from tinytensor.models import LeNet
from tinytensor.optim  import AdamW
from tinytensor.nn     import CrossEntropyLoss

model = LeNet(num_classes=10, in_channels=1)
model.compile(lambda p: AdamW(p, lr=1e-3), CrossEntropyLoss())

model.fit(x_train, y_train, epochs=10,
          validation_data=(x_test, y_test))

acc = (model.predict(x_test) == y_test).mean()

возможности

Всё для обучения и сжатия

От autograd до развёртывания на edge-устройствах.

∇

Autograd

Backprop по графу. Пишешь loss.backward() — градиенты текут сами.

▤

Свёртки

Conv2d, BatchNorm, пулинг, ResidualBlock. Готовые LeNet, ResNet, VGG.

◧

Квантизация INT8

Симметричная per-channel. Сжатие вчетверо вызовом .quant().

✂

Прунинг

Режет нейроны по важности. Реально уменьшает FLOPs, а не зануляет веса.

⊞

Ансамбли

Несколько моделей, hard/soft голосование. Один Ensemble([...]).

◑

Колбэки

Чекпоинты, ранняя остановка, csv-логи, свои хуки. Как в Keras.

⤢

ONNX

Экспорт и импорт. Открывается в Netron и ONNX Runtime.

◴

Профилирование

benchmark, summary, FLOPs, размер. Под замеры на Raspberry Pi.

⚡

CUDA

Один tt.set_device("cuda") — весь код идёт на GPU.

как работает

Граф, по которому течёт градиент

Каждая операция помнит, как посчитать свою производную. Тензор держит родителей и функцию градиента. Вызов .backward() идёт по графу назад и раздаёт градиенты. Свой слой — это Module с методом forward, backward считается сам.

autograd.py

# y = x² + 2x,  dy/dx = 2x + 2
x = Tensor([3.0], requires_grad=True)
y = x*x + Tensor([2.0])*x

y.backward()
x.grad   # [8.]