A fully GPU-accelerated neural-network training pipeline implemented in **WebGPU** using custom **WGSL compute shaders**.
This project trains a **2-layer MLP classifier** on MNIST entirely inside the browser with:
* Forward pass (dense layer 1 + ReLU)
* Forward pass (dense layer 2 → logits)
* Softmax + cross-entropy loss
* Backpropagation for both dense layers
* Weight/bias gradient computation
* SGD optimizer update
* Live accuracy + loss visualization
* Browser-side image preview of predictions
No ML frameworks and no libraries — every compute pass, buffer, math operation, and gradient update runs through raw WebGPU shaders.
---
## Features
### ✔ Fully WebGPU-accelerated training
All computation (forward pass, backward pass, loss, optimizer) is written in WGSL compute kernels.
### ✔ End-to-end MNIST pipeline
Images and labels are uploaded to GPU buffers and remain on GPU throughout training.
### ✔ Pure WGSL neural network
Six compute shaders:
1.`forward_dense_layer1.wgsl`
2.`forward_dense_layer2.wgsl`
3.`softmax_cross_entropy_backward.wgsl`
4.`gradients_dense_layer2.wgsl`
5.`gradients_dense_layer1.wgsl`
6.`sgd_optimizer_update.wgsl`
### ✔ Pure JS orchestration
A lightweight JS engine (`training2.js`) handles:
* shader setup
* buffer binding
* dispatch sizes
* ordered forward/backward passes
* debug readbacks
* canvas preview of sample predictions
### ✔ Modular WebGPU runtime
The training pipeline uses a generic WebGPU shader wrapper (`WebGpu.js`) that automatically:
* parses WGSL bindings
* allocates buffers
* recreates bind groups
* handles resizing
* dispatches compute passes
* reads debug buffers
---
## Live Results
A typical run with batch size **64**:
```
Epoch 0: Loss = 2.3032 Accuracy = 9.40%
Epoch 1: Loss = 2.1380 Accuracy = 26.40%
Epoch 5: Loss = 1.8169 Accuracy = 36.90%
Epoch 10: Loss = 1.6810 Accuracy = 40.60%
Epoch 19: Loss = 1.6004 Accuracy = 41.60%
```
This matches what you'd expect from a small, unoptimized 128-unit MLP trained only on a single batch — validating that all gradients and updates propagate correctly.
