Howdy,
I recently finished my undergrad and wanted to learn more about Neural Networks, Mojo , and GPU kernels, so I decided to write LeNet-5* from scratch without any additional libraries for the MNIST data set.
Models can train and test on CPU, and test on GPU (no GPU training for now). You can also load a trained model for testing from a previous project of mine. Performance on GPU (with a batch size of 50) appears to be 4% faster than a PyTorch implementation, which I’m happy with. There are also some simple logging features available.
I hope it’s a fun example of what Mojo is capable of, how a CNN works, and how approachable writing GPU kernels can be.
I’d love any comments or feedback!
Cheers!
* There are some differences between this version and the one described in the original LeCun et al paper, as noted in the project.