Spaces:
Sleeping
Sleeping
| title: TinyTorch Viz | |
| emoji: 🔥 | |
| colorFrom: blue | |
| colorTo: purple | |
| sdk: docker | |
| app_port: 7860 | |
| pinned: false | |
| # TinyTorch Viz 🔥 | |
| A lightweight, educational framework built to visualize your deep learning architectures and foster a deep mathematical intuition. | |
| Based on [TinyTorch](https://mlsysbook.ai/tinytorch/intro.html) and the educational concepts of [AI by Hand](https://www.byhand.ai/), making it perfect for learning ML fundamentals ahead of the curve. | |
| ## 🚀 Quick Start | |
| **Try it instantly on Hugging Face Spaces:** | |
| [**Launch TinyTorch Viz**](https://huggingface.co/spaces/gabriead/tiny-torch-viz) | |
| ### Or run locally: | |
| ```bash | |
| cd /path/to/TinyTorch | |
| uv sync | |
| source .venv/bin/activate | |
| uv run uvicorn app:app --host 0.0.0.0 --port 8000 | |
| ``` | |
| Open http://localhost:8000 in your browser. | |
| --- | |
| ## ✨ Features | |
| - **Visual Code Execution:** Write TinyTorch code and instantly generate a dynamic visual representation of your tensors and operations. This connects abstract code to concrete visual structures, deepening your understanding of neural network mechanics. | |
| - **Build Architectures from Scratch:** Create complete Deep Learning models layer-by-layer, visualizing the flow of data and gradients through every component. | |
| - **Scalable Visualization:** Toggle the "Data" view to hide raw tensor values and focus purely on architecture. This allows you to design larger, complex networks without clutter. | |
| - **Side-by-Side PDF Learning:** Open research papers or tutorials directly within the app to code alongside the text, enabling immediate implementation and verification of concepts. | |
| - **Rich Annotations:** Add notes with full LaTeX support to explain mathematical formulas, document layer logic, or leave reminders directly on your visual representation. | |
| - **Save & Share:** Export your entire workspace—including code, visual representation, and notes—as a JSON file to share your educational content or resume work later. | |
| - **Retro Mode:** Switch to a distinct retro-styled UI for a focused, nostalgic coding experience. | |
| ## ⚠️ Important Note on Math Notation | |
| **I changed the order of matrix multiplication from $x \cdot W$ to $W \cdot x$ in the Linear layer following the learning concept of Prof. Tom Yeh.** | |
| This follows the [AI by Hand](https://www.byhand.ai/) teaching methodology. While standard libraries often use `input @ weights`, this tool visualizes weights first to align with standard mathematical notation found in linear algebra textbooks. | |
| --- | |
| ## 💡 Ideas to get started | |
| - **Visualize Basic Operations:** Start simple by creating two tensors and visualizing their addition or matrix multiplication. This helps you understand how `box()` commands group operations in the UI. | |
| - **Recreate a Textbook Example:** Take a small neural network diagram from a textbook or paper and implement it layer-by-layer to see if the data flow matches your expectation. | |
| - **Debug a "Shape Mismatch":** Intentionally create a dimension error (like multiplying `(4,8)` by `(4,16)`) to see how the visualizer handles it versus standard Python error messages. | |
| - **Annotate Your Math:** Use the LaTeX note feature to write down the formula $y = \sigma(Wx + b)$ next to your Linear layer visualization to connect the code to the math. | |
| ## 📚 Core API Reference | |
| Please look at the great work done by [TinyTorch](https://mlsysbook.ai/tinytorch/intro.html) for a complete API reference. Below are the currently enabled features for visualization. | |
| The examples are all taken freom the TinyTorch API. | |
| ### 📦 Tensor | |
| The fundamental data structure. | |
| Python | |
| ``` | |
| import numpy as np | |
| # Create tensors | |
| a = Tensor([1, 2, 3, 4]) | |
| b = Tensor(np.random.randn(3, 4)) | |
| ``` | |
| #### Arithmetic Operations | |
| Python | |
| ``` | |
| c = a + b # Addition | |
| d = a - b # Subtraction | |
| e = a * b # Multiplication | |
| f = a / b # Division | |
| ``` | |
| #### Matrix Operations | |
| Python | |
| ``` | |
| a = Tensor([[1, 2, 3], [4, 5, 6]]) # Shape: (2, 3) | |
| b = Tensor([[1, 2], [3, 4], [5, 6]]) # Shape: (3, 2) | |
| # Matrix multiplication | |
| c = a @ b # Using @ operator | |
| c = a.matmul(b) # Using method | |
| # Transpose | |
| d = a.transpose() # Swap last two dimensions | |
| ``` | |
| #### Reduction Operations | |
| Python | |
| ``` | |
| total = a.sum() # Sum all elements | |
| col_sums = a.sum(axis=0) # Sum each column | |
| avg = a.mean() # Mean of all elements | |
| maximum = a.max() # Max of all elements | |
| ``` | |
| ### ⚡ Activations | |
| Non-linear functions that enable neural networks to learn complex patterns. | |
| - `ReLU()` | |
| - `Sigmoid()` | |
| - `Tanh()` | |
| - `GELU()` | |
| - `Softmax()` | |
| ### 🧱 Layers | |
| Building blocks for neural network architectures. | |
| #### Linear Layer | |
| Fully connected layer: $y = Wx + b$ (Note the order!) | |
| Python | |
| ``` | |
| # Create layer: 8 input features → 4 output features | |
| linear = Linear(8, 4) | |
| y = linear(x) | |
| ``` | |
| #### Dropout | |
| Regularization layer. | |
| Python | |
| ``` | |
| dropout = Dropout(p=0.5) | |
| dropout.train() # Drop values | |
| y = dropout(x) | |
| ``` | |
| #### Sequential | |
| Container for stacking layers. | |
| Python | |
| ``` | |
| model = Sequential( | |
| Linear(784, 256), | |
| ReLU(), | |
| Linear(256, 10), | |
| Softmax() | |
| ) | |
| ``` | |
| ### 📉 Loss Functions | |
| - `MSELoss()`: For regression tasks. | |
| - `CrossEntropyLoss()`: For multi-class classification. | |
| ### 🎨 Visualization Commands | |
| The visualizer provides special commands to organize the UI. | |
| Python | |
| ``` | |
| # Group tensors into a colored box | |
| box("Layer 1", [x, weights, output], "2") | |
| ``` | |
| **Color Schemes:** | |
| - `"1"` - Default (dark) | |
| - `"2"` - Green (ReLU/activation) | |
| - `"3"` - Blue (Linear layers) | |
| - `"4"` - Purple (Softmax) | |
| - `"5"` - Orange (Sigmoid) | |
| - `"6"` - Red (Dropout/Loss) | |
| --- | |
| ## 🧪 Complete Example: 3-Layer Neural Network | |
| Python | |
| ``` | |
| import numpy as np | |
| # 1. Create Data (Note shapes: Batch=4, Features=8) | |
| X = Tensor(np.random.randn(4, 8)) | |
| y = Tensor([0, 1, 2, 1]) | |
| # 2. Define Network | |
| layer1 = Linear(8, 16) | |
| relu1 = ReLU() | |
| layer2 = Linear(16, 8) | |
| relu2 = ReLU() | |
| layer3 = Linear(8, 3) | |
| softmax = Softmax() | |
| loss_fn = CrossEntropyLoss() | |
| # 3. Forward Pass | |
| z1 = layer1(X) | |
| a1 = relu1(z1) | |
| z2 = layer2(a1) | |
| a2 = relu2(z2) | |
| logits = layer3(a2) | |
| predictions = softmax(logits) | |
| # 4. Compute Loss | |
| loss = loss_fn(logits, y) | |
| # 5. Visualize | |
| box("Input", X, "1") | |
| box("Layer 1", [z1, a1], "2") | |
| box("Output", [logits, predictions, loss], "4") | |
| ``` | |
| --- | |
| ## 📁 Project Structure | |
| ``` | |
| TinyTorch/ | |
| ├── app.py # FastAPI server for visualization | |
| ├── tracer.py # Tensor operation tracing | |
| ├── instrumentation.py # Hooks for tracing | |
| ├── static/ | |
| │ └── index.html # Visualization frontend | |
| └── tinytorch/ | |
| └── core/ # Core ML library implementation | |
| ``` | |
| ## License | |
| MIT License - Built for education and understanding. | |