Building a Deep learning framework in C++ (from scratch) - training MNIST as a milestone

i am building a deep learning framework called "Forge" completely from scratch in C++, its nowhere near complete yet, training MNIST Classifier shows a functional core on CPU (i'll add a CUDA backend too). My end goal is to train a modern transformer on Forge.

YT video of MNIST training :- www.youtube.com/watch?v=CalrXYYmpfc

this video shows:

\-> training an MLP on MNIST
\-> loss decreasing over epochs
\-> predictions vs ground truth

this stable training proves that the following components are working correctly:-

\--> Tensor system (it uses Eigen as math backend, but i'll handcraft the math backend/kernels for CUDA later) and CPU memory allocator.
\--> autodiff engine (computation graph is being built and traversed correctly)
\-->primitives -- linear layer, relu activation (Forge has sigmoid, softmax, gelu, tanh and leakyrelu too), CrossEntropy loss function (it fuses log softmax and CE. Forge has MSE and BinaryCrossEntropy too, the BCE fuses sigmoid and BCE) and SGD optimizer (i am planning to add momentum in SGD, Adam and AdamW)

[the Forge repo on GitHub is currently private as its WAP\]
My GitHub: github.com/muchlakshay

https://redd.it/1silamh
@r_cpp