MNIST 3D MLP Visualization

This project trains a tiny MLP on MNIST with PyTorch and visualizes it in 3D using PyOpenGL + GLFW. The scene matches the chaotic 3D MLP style:

• 28×28 input pixels as a floating grid of squares.
• Hidden neurons rendered as glowing particles in 3D, positioned by a force-directed layout.
• Semi-transparent edges between neurons; brightness tracks weight magnitude (positive vs negative hue).
• Camera pan/rotate/zoom controls for free-flight exploration.

Quick start

python -m venv .venv
source .venv/bin/activate  # Windows: .venv\\Scripts\\activate
pip install -r requirements.txt
python main.py

Notes:

• The first run downloads MNIST and trains for a single epoch; it caches the model at checkpoints/mlp_mnist.pt. Re-run to skip training.
• Use --epochs N to train longer, --no-train to only load the checkpoint, and --sample-index K to visualize a specific test image.
• If you want to limit edges for performance, set --max-edges 8000 (default) or a lower value.

Controls

• Left mouse drag: orbit camera.
• Right mouse drag: pan camera.
• Scroll: zoom in/out.
• Space: reseed force layout (reshuffle hidden clusters).
• N: pick a new random MNIST test sample immediately (auto-rotates every 2s).
• Esc / Q: quit.

Files

• main.py – entrypoint; loads/trains the network, grabs a sample, builds the scene, and runs the render loop.
• mlp.py – simple PyTorch MLP (784 → 256 → 128 → 10) plus activation capture.
• train.py – minimal MNIST training and checkpointing.
• layout.py – force-directed layout for hidden/output neurons and fixed grid for inputs.
• renderer.py – PyOpenGL + GLFW renderer, camera controls, and scene drawing.

Troubleshooting

• If GLFW fails to init, ensure system OpenGL drivers are installed.
• For headless servers, use a virtual display (e.g., Xvfb).
• To debug rendering, run with --wireframe to see the geometry without blending.***