mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-03-12 01:23:34 -05:00
768aea4aa9
MAJOR FEATURE: Multi-channel convolutions for real CNN architectures Key additions: - MultiChannelConv2D class with in_channels/out_channels support - Handles RGB images (3 channels) and arbitrary channel counts - He initialization for stable training - Optional bias parameters - Batch processing support Testing & Validation: - Comprehensive unit tests for single/multi-channel - Integration tests for complete CNN pipelines - Memory profiling and parameter scaling analysis - QA approved: All mandatory tests passing CIFAR-10 CNN Example: - Updated train_cnn.py to use MultiChannelConv2D - Architecture: Conv(3→32) → Pool → Conv(32→64) → Pool → Dense - Demonstrates why convolutions matter for vision - Shows parameter reduction vs MLPs (18KB vs 12MB) Systems Analysis: - Parameter scaling: O(in_channels × out_channels × kernel²) - Memory profiling shows efficient scaling - Performance characteristics documented - Production context with PyTorch comparisons This enables proper CNN training on CIFAR-10 with ~60% accuracy target.
4.7 KiB
4.7 KiB
TinyTorch Module Audit: Essential vs Extra Components
Overview
This audit examines what components are NEEDED for each milestone vs EXTRA components that enhance the framework but aren't strictly necessary.
Part I: MLPs (Target: XORNet)
Module 02: Tensor
ESSENTIAL for XORNet:
- Basic Tensor class with data storage
- Addition, subtraction, multiplication
- Matrix multiply (for layers)
- Shape, reshape operations
EXTRA (but good for framework):
- Broadcasting ✓ (nice but XOR doesn't need)
- Fancy indexing ✓
- Statistical operations (mean, sum, std) ✓
- Comparison operators ✓
Module 03: Activations
ESSENTIAL for XORNet:
- ReLU ✓ (used in XORNet)
- Sigmoid (could use for XOR output)
EXTRA (but good for framework):
- Tanh ✓ (alternative to ReLU)
- Softmax ✓ (not needed for XOR, but needed for CIFAR-10)
- ActivationProfiler ✓ (pedagogical tool)
Module 04: Layers
ESSENTIAL for XORNet:
- Dense layer ✓ (fully connected)
- Weight initialization
- Forward pass
EXTRA:
- Different initialization strategies (Xavier, He, etc.)
- Bias option
Module 05: Networks
ESSENTIAL for XORNet:
- Sequential model ✓
- Forward pass through layers
EXTRA:
- Model summary/printing
- Parameter counting
Part II: CNNs (Target: CIFAR-10)
Module 06: Spatial
ESSENTIAL for CNN CIFAR-10:
- Conv2D ✓ (the key innovation!)
- MaxPool2D ✓ (for downsampling)
EXTRA (but pedagogically valuable):
- Different padding modes
- Stride options
- AvgPool2D (alternative pooling)
- Multiple filter support
Module 07: DataLoader
ESSENTIAL for CIFAR-10:
- CIFAR10Dataset ✓
- DataLoader with batching ✓
- Shuffling ✓
EXTRA:
- Data augmentation (but helps accuracy!)
- Other datasets (MNIST, etc.)
- Prefetching/parallel loading
Module 08: Autograd
ESSENTIAL for CIFAR-10:
- Variable class ✓
- Backward pass ✓
- Gradient computation ✓
EXTRA:
- Computation graph visualization
- Gradient checking
- Higher-order derivatives
Module 09: Optimizers
ESSENTIAL for CIFAR-10:
- SGD (basic, could work)
- Adam ✓ (used in CIFAR-10, converges faster)
EXTRA:
- Learning rate scheduling
- Momentum variants
- RMSprop, AdaGrad
Module 10: Training
ESSENTIAL for CIFAR-10:
- Training loop ✓
- CrossEntropyLoss ✓
- Basic evaluation ✓
EXTRA (but very useful):
- Checkpointing ✓
- Early stopping ✓
- Metrics tracking ✓
- Validation splits ✓
- MeanSquaredError (for XOR)
Part III: Transformers (Target: TinyGPT)
Module 11: Embeddings
ESSENTIAL for TinyGPT:
- Token embedding layer
- Positional encoding (sinusoidal or learned)
EXTRA:
- Multiple embedding types
- Embedding dropout
Module 12: Attention
ESSENTIAL for TinyGPT:
- Multi-head attention ✓ (already implemented!)
- Scaled dot-product attention ✓
- Causal masking ✓
EXTRA:
- Different attention variants
- Attention visualization
Module 13: Normalization
ESSENTIAL for TinyGPT:
- LayerNorm (critical for transformer stability)
EXTRA:
- BatchNorm (not used in transformers)
- GroupNorm, InstanceNorm
Module 14: Transformers
ESSENTIAL for TinyGPT:
- TransformerBlock (attention + FFN + residual)
- Positional encoding integration
- Stack of blocks
EXTRA:
- Encoder-decoder architecture
- Cross-attention
Module 15: Generation
ESSENTIAL for TinyGPT:
- Autoregressive generation
- Temperature sampling
- Greedy decoding
EXTRA:
- Beam search
- Top-k, Top-p sampling
- Repetition penalty
Summary
Truly Minimal Path
If we wanted ONLY what's needed for milestones:
- XORNet: Just needs Dense, ReLU, basic Tensor ops
- CIFAR-10 MLP: Add DataLoader, Adam, CrossEntropyLoss
- CIFAR-10 CNN: Add Conv2D, MaxPool2D
- TinyGPT: Add Embeddings, Attention, LayerNorm, Generation
What We Have (Good Extras)
- More activation choices: Good for experimentation
- Better optimizers: Adam converges faster than SGD
- Training utilities: Checkpointing, metrics (very practical!)
- Profiling tools: Help understand performance
Missing Essentials
For Part III (TinyGPT) we still need to implement:
- Module 11: Embedding layer, positional encoding
- Module 13: LayerNorm
- Module 14: TransformerBlock
- Module 15: Generation strategies
Verdict
The current modules have a good balance of essential + useful extras. The extras are:
- Pedagogically valuable (show alternatives)
- Practically useful (checkpointing, better optimizers)
- Framework completeness (makes TinyTorch feel real)
The only "bloat" might be multiple activation functions, but even those are good for showing students the options and tradeoffs.