I unfortunately didn’t get to finish GPU kernels this weekend, but here this is. CPU only forward passes work with 97.88% accuracy on 10,000 MNIST images using a pretrained model.
I chose this because in a class in my last semester we worked on a C / CUDA version that taught me a lot and I figured it was good to keep digging in and especially focus on understanding backprop this time and then optimizing GPU kernels.
I wanted to get further this weekend but life happened, GPU kernels will be next along with writing backprop for training in Mojo and digging into SIMD support. Then I’m hoping to do quantization and pruning afterwards. I also am certain I can organize the code better since Mojo has features C kinda doesn’t (struct methods), I wanted a close one-to-one of the source just for reference.
I’m a recent undergrad, so a lot of my problems I’m sure are mostly due to that. Some of it was certainly skill issues or not reading things as well as I should’ve (doing .stack_allocation() on a LayoutTensor means it goes out of scope really fast, duh. I still don’t know how to allocate my LeNet5 struct on the stack).
InlineArrays were giving me some trouble near the end (alignment / span couldn’t be inferred when passing to a LayoutTensor as storage) so I swapped to UnsafePointers. Other things like this happened which are probably just skill issues. Loading the model was tedious. The language is complex and it’ll take time and projects and probably reading more of the stdlib to help me get further.
I’d really only done a few days in Advent of Code before this and ~8 of the GPU puzzles so this was a new language to me largely.