mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-07-16 12:22:52 -05:00
7d41bb125e
- Remove top-level SimpleModel from modules 15 & 16 (keep in test functions) - Rename QuantizationComplete → Quantizer (cleaner, matches Profiler pattern) - Rename CompressionComplete → Compressor (same pattern) - Rename benchmarking.benchmark → bench (shorter)
523 lines
17 KiB
Python
523 lines
17 KiB
Python
"""
|
|
Module 05: Autograd - Progressive Testing
|
|
==========================================
|
|
|
|
🎯 LEARNING OBJECTIVES:
|
|
1. Understand automatic differentiation
|
|
2. Build computation graphs during forward pass
|
|
3. Compute gradients via backpropagation
|
|
|
|
📚 PREREQUISITE MODULES:
|
|
- Module 01: Tensor (data structure)
|
|
- Module 02: Activations (non-linear functions)
|
|
- Module 03: Layers (Linear transformation)
|
|
- Module 04: Losses (objective functions)
|
|
|
|
🔗 WHAT AUTOGRAD ENABLES:
|
|
After this module, your tensors can automatically compute gradients!
|
|
This is the foundation of neural network training.
|
|
"""
|
|
|
|
import pytest
|
|
import numpy as np
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add project root
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
|
|
|
|
|
# =============================================================================
|
|
# SECTION 1: REGRESSION TESTS
|
|
# Verify earlier modules still work after autograd patches tensors
|
|
# =============================================================================
|
|
|
|
class TestFoundationStillWorks:
|
|
"""
|
|
🛡️ REGRESSION CHECK: Autograd must not break the foundation
|
|
|
|
Autograd patches Tensor operations to track gradients. This test ensures
|
|
basic tensor functionality still works correctly after enabling autograd.
|
|
|
|
WHY THIS MATTERS:
|
|
A common bug is breaking basic operations when adding gradient tracking.
|
|
If tensor creation or arithmetic breaks, nothing else will work!
|
|
"""
|
|
|
|
def test_tensor_creation_works(self):
|
|
"""
|
|
✅ WHAT: Basic tensor creation
|
|
🔍 IF FAILS: Autograd broke the Tensor constructor
|
|
"""
|
|
from tinytorch import Tensor
|
|
|
|
# These should all still work
|
|
t1 = Tensor([1, 2, 3])
|
|
t2 = Tensor([[1, 2], [3, 4]])
|
|
t3 = Tensor(np.random.randn(3, 4, 5))
|
|
|
|
assert t1.shape == (3,), "1D tensor creation broken"
|
|
assert t2.shape == (2, 2), "2D tensor creation broken"
|
|
assert t3.shape == (3, 4, 5), "3D tensor creation broken"
|
|
|
|
def test_tensor_arithmetic_works(self):
|
|
"""
|
|
✅ WHAT: Basic arithmetic (+, -, *, /)
|
|
🔍 IF FAILS: Autograd broke tensor operators
|
|
"""
|
|
from tinytorch import Tensor
|
|
|
|
a = Tensor([1.0, 2.0, 3.0])
|
|
b = Tensor([4.0, 5.0, 6.0])
|
|
|
|
# All basic operations should work
|
|
add_result = a + b
|
|
sub_result = a - b
|
|
mul_result = a * b
|
|
div_result = a / b
|
|
|
|
assert np.allclose(add_result.data, [5, 7, 9]), "Addition broken"
|
|
assert np.allclose(sub_result.data, [-3, -3, -3]), "Subtraction broken"
|
|
assert np.allclose(mul_result.data, [4, 10, 18]), "Multiplication broken"
|
|
assert np.allclose(div_result.data, [0.25, 0.4, 0.5]), "Division broken"
|
|
|
|
def test_linear_layer_still_works(self):
|
|
"""
|
|
✅ WHAT: Linear layer forward pass
|
|
🔍 IF FAILS: Autograd broke layer operations
|
|
"""
|
|
from tinytorch import Tensor, Linear
|
|
|
|
layer = Linear(10, 5)
|
|
x = Tensor(np.random.randn(3, 10)) # batch of 3
|
|
|
|
output = layer(x)
|
|
|
|
assert output.shape == (3, 5), (
|
|
f"Linear layer output shape wrong!\n"
|
|
f" Input: (3, 10)\n"
|
|
f" Expected output: (3, 5)\n"
|
|
f" Got: {output.shape}\n"
|
|
f"\n"
|
|
f"💡 HINT: Linear(10, 5) should transform (batch, 10) → (batch, 5)"
|
|
)
|
|
|
|
|
|
class TestActivationsStillWork:
|
|
"""
|
|
🛡️ REGRESSION CHECK: Activations must still work with autograd-enabled tensors
|
|
"""
|
|
|
|
def test_relu_works_with_gradients(self):
|
|
"""
|
|
✅ WHAT: ReLU on tensors that require gradients
|
|
🔍 IF FAILS: ReLU doesn't handle requires_grad properly
|
|
"""
|
|
from tinytorch import Tensor, ReLU
|
|
|
|
relu = ReLU()
|
|
x = Tensor([-2, -1, 0, 1, 2], requires_grad=True)
|
|
|
|
output = relu(x)
|
|
|
|
assert np.allclose(output.data, [0, 0, 0, 1, 2]), (
|
|
"ReLU computation wrong!\n"
|
|
" Input: [-2, -1, 0, 1, 2]\n"
|
|
" Expected: [0, 0, 0, 1, 2]\n"
|
|
f" Got: {output.data}\n"
|
|
"\n"
|
|
"💡 HINT: ReLU(x) = max(0, x)"
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# SECTION 2: CAPABILITY TESTS
|
|
# Verify Module 05 provides its core functionality
|
|
# =============================================================================
|
|
|
|
class TestAutogradCapabilities:
|
|
"""
|
|
🎯 CAPABILITY CHECK: Does autograd do what it's supposed to?
|
|
|
|
Autograd must:
|
|
1. Track operations during forward pass (build computation graph)
|
|
2. Compute gradients during backward pass (backpropagation)
|
|
3. Store gradients in .grad attribute
|
|
"""
|
|
|
|
def test_requires_grad_flag_exists(self):
|
|
"""
|
|
✅ WHAT: Tensors have requires_grad attribute
|
|
|
|
📖 CONCEPT: requires_grad tells autograd whether to track this tensor
|
|
- requires_grad=True → track operations, compute gradients
|
|
- requires_grad=False → don't track (saves memory)
|
|
"""
|
|
from tinytorch import Tensor
|
|
|
|
t1 = Tensor([1, 2, 3], requires_grad=True)
|
|
t2 = Tensor([1, 2, 3], requires_grad=False)
|
|
t3 = Tensor([1, 2, 3]) # default
|
|
|
|
assert hasattr(t1, 'requires_grad'), "Tensor missing requires_grad attribute"
|
|
assert t1.requires_grad == True, "requires_grad=True not stored"
|
|
assert t2.requires_grad == False, "requires_grad=False not stored"
|
|
|
|
def test_grad_attribute_exists(self):
|
|
"""
|
|
✅ WHAT: Tensors have .grad attribute for storing gradients
|
|
|
|
📖 CONCEPT: After backward(), gradients are stored in .grad
|
|
"""
|
|
from tinytorch import Tensor
|
|
|
|
t = Tensor([1, 2, 3], requires_grad=True)
|
|
|
|
assert hasattr(t, 'grad'), (
|
|
"Tensor missing .grad attribute!\n"
|
|
"\n"
|
|
"💡 HINT: Add 'self.grad = None' in Tensor.__init__()"
|
|
)
|
|
|
|
def test_simple_gradient_computation(self):
|
|
"""
|
|
✅ WHAT: Gradients computed for y = sum(x * 2)
|
|
|
|
📖 CONCEPT: If y = sum(2x), then dy/dx = 2 for each element
|
|
We use sum() to get a scalar for backward().
|
|
|
|
🔍 IF FAILS: Your backward pass isn't working
|
|
"""
|
|
from tinytorch import Tensor
|
|
|
|
x = Tensor([1.0, 2.0, 3.0], requires_grad=True)
|
|
y = x * 2 # Simple operation
|
|
loss = y.sum() # Must be scalar for backward()
|
|
|
|
# Backward pass
|
|
loss.backward()
|
|
|
|
assert x.grad is not None, (
|
|
"Gradient not computed!\n"
|
|
"\n"
|
|
"For y = 2x, we expect dy/dx = 2\n"
|
|
"\n"
|
|
"💡 HINTS:\n"
|
|
"1. Is backward() calling the right backward function?\n"
|
|
"2. Are gradients being stored in .grad?"
|
|
)
|
|
|
|
expected_grad = np.array([2.0, 2.0, 2.0])
|
|
assert np.allclose(x.grad, expected_grad), (
|
|
f"Gradient value wrong!\n"
|
|
f" For y = 2x, dy/dx should be 2\n"
|
|
f" Expected: {expected_grad}\n"
|
|
f" Got: {x.grad}\n"
|
|
f"\n"
|
|
"💡 HINT: Check your multiplication backward function"
|
|
)
|
|
|
|
def test_chain_rule_works(self):
|
|
"""
|
|
✅ WHAT: Gradients flow through multiple operations (chain rule)
|
|
|
|
📖 CONCEPT: Chain Rule
|
|
If z = g(y) and y = f(x), then:
|
|
dz/dx = dz/dy * dy/dx
|
|
|
|
This is the foundation of backpropagation!
|
|
|
|
Example: loss = sum((x * 2) + 3)
|
|
- y = x * 2 → dy/dx = 2
|
|
- z = y + 3 → dz/dy = 1
|
|
- loss = sum(z) → dloss/dz = 1
|
|
- Therefore: dloss/dx = 1 * 1 * 2 = 2
|
|
"""
|
|
from tinytorch import Tensor
|
|
|
|
x = Tensor([1.0, 2.0, 3.0], requires_grad=True)
|
|
y = x * 2 # dy/dx = 2
|
|
z = y + 3 # dz/dy = 1
|
|
loss = z.sum() # Must be scalar for backward()
|
|
|
|
loss.backward()
|
|
|
|
expected_grad = np.array([2.0, 2.0, 2.0]) # dz/dx = 2
|
|
assert x.grad is not None, "Chain rule: gradients didn't flow back"
|
|
assert np.allclose(x.grad, expected_grad), (
|
|
f"Chain rule gradient wrong!\n"
|
|
f" z = (x * 2) + 3\n"
|
|
f" dz/dx = dz/dy * dy/dx = 1 * 2 = 2\n"
|
|
f" Expected: {expected_grad}\n"
|
|
f" Got: {x.grad}"
|
|
)
|
|
|
|
|
|
class TestNeuralNetworkGradients:
|
|
"""
|
|
🎯 CAPABILITY CHECK: Can autograd train neural networks?
|
|
|
|
This is the real test: can we compute gradients for a neural network?
|
|
"""
|
|
|
|
def test_linear_layer_gradients(self):
|
|
"""
|
|
✅ WHAT: Gradients flow through Linear layer
|
|
|
|
📖 CONCEPT: For y = xW + b:
|
|
- dy/dW = x^T (input transposed)
|
|
- dy/db = 1 (gradient of bias is 1)
|
|
- dy/dx = W^T (weight transposed)
|
|
"""
|
|
from tinytorch import Tensor, Linear
|
|
|
|
# Simple linear layer
|
|
layer = Linear(3, 2)
|
|
x = Tensor([[1.0, 2.0, 3.0]], requires_grad=True)
|
|
|
|
# Forward
|
|
y = layer(x)
|
|
|
|
# Create simple loss (sum of outputs)
|
|
loss = y.sum()
|
|
|
|
# Backward
|
|
loss.backward()
|
|
|
|
# Weight should have gradients
|
|
assert layer.weight.grad is not None, (
|
|
"Linear layer weights didn't receive gradients!\n"
|
|
"\n"
|
|
"💡 HINTS:\n"
|
|
"1. Is layer.weight.requires_grad = True?\n"
|
|
"2. Did you implement matmul backward correctly?\n"
|
|
"3. Are gradients propagating through the add operation?"
|
|
)
|
|
|
|
# Bias should have gradients
|
|
if layer.bias is not None:
|
|
assert layer.bias.grad is not None, (
|
|
"Linear layer bias didn't receive gradients!"
|
|
)
|
|
|
|
def test_mlp_end_to_end_gradients(self):
|
|
"""
|
|
✅ WHAT: Multi-layer network computes gradients
|
|
|
|
📖 CONCEPT: Backprop through multiple layers
|
|
Each layer receives gradients from the layer above.
|
|
"""
|
|
from tinytorch import Tensor, Linear, ReLU
|
|
|
|
# Two-layer MLP
|
|
layer1 = Linear(4, 8)
|
|
relu = ReLU()
|
|
layer2 = Linear(8, 2)
|
|
|
|
# Forward
|
|
x = Tensor(np.random.randn(2, 4), requires_grad=True)
|
|
h = layer1(x)
|
|
h = relu(h)
|
|
y = layer2(h)
|
|
|
|
# Loss and backward
|
|
loss = y.sum()
|
|
loss.backward()
|
|
|
|
# All layers should have gradients
|
|
assert layer1.weight.grad is not None, "Layer 1 didn't receive gradients"
|
|
assert layer2.weight.grad is not None, "Layer 2 didn't receive gradients"
|
|
|
|
# Gradients should be non-zero
|
|
assert np.any(layer1.weight.grad != 0), (
|
|
"Layer 1 has zero gradients!\n"
|
|
"\n"
|
|
"💡 HINT: Check if gradients are flowing through ReLU.\n"
|
|
"ReLU gradient is 1 for positive inputs, 0 for negative."
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# SECTION 3: INTEGRATION TESTS
|
|
# Verify autograd works with all previous modules together
|
|
# =============================================================================
|
|
|
|
class TestAutogradLossIntegration:
|
|
"""
|
|
🔗 INTEGRATION CHECK: Autograd + Loss functions
|
|
|
|
Training requires computing gradients of the loss.
|
|
"""
|
|
|
|
def test_mse_loss_gradients(self):
|
|
"""
|
|
✅ WHAT: MSE loss produces correct gradients
|
|
|
|
📖 CONCEPT: MSE = mean((predictions - targets)^2)
|
|
Gradient: d(MSE)/d(predictions) = 2 * (predictions - targets) / n
|
|
"""
|
|
from tinytorch import Tensor, MSELoss
|
|
|
|
predictions = Tensor([[1.0, 2.0, 3.0]], requires_grad=True)
|
|
targets = Tensor([[1.5, 2.5, 2.5]])
|
|
|
|
loss_fn = MSELoss()
|
|
loss = loss_fn(predictions, targets)
|
|
|
|
loss.backward()
|
|
|
|
assert predictions.grad is not None, (
|
|
"MSE loss didn't produce gradients!\n"
|
|
"\n"
|
|
"💡 HINT: Is loss.backward() calling the right backward function?"
|
|
)
|
|
|
|
|
|
class TestCompleteTrainingLoop:
|
|
"""
|
|
🔗 INTEGRATION CHECK: Can we do one complete training step?
|
|
|
|
This tests everything together:
|
|
1. Forward pass through layers
|
|
2. Compute loss
|
|
3. Backward pass (autograd)
|
|
4. Verify gradients exist for optimization
|
|
"""
|
|
|
|
def test_training_step_computes_gradients(self):
|
|
"""
|
|
✅ WHAT: Complete forward-backward pass works
|
|
|
|
This is what happens in every training step:
|
|
1. Feed data through network
|
|
2. Compute loss
|
|
3. Compute gradients
|
|
4. (Optimizer would update weights here)
|
|
"""
|
|
from tinytorch import Tensor, Linear, ReLU, MSELoss
|
|
|
|
# Simple network
|
|
layer = Linear(4, 2)
|
|
activation = ReLU()
|
|
|
|
# Data
|
|
x = Tensor(np.random.randn(8, 4)) # 8 samples
|
|
target = Tensor(np.random.randn(8, 2))
|
|
|
|
# Forward
|
|
hidden = layer(x)
|
|
output = activation(hidden)
|
|
|
|
# Loss
|
|
loss_fn = MSELoss()
|
|
loss = loss_fn(output, target)
|
|
|
|
# Backward
|
|
loss.backward()
|
|
|
|
# Verify gradients exist
|
|
assert layer.weight.grad is not None, (
|
|
"Training step failed: weights have no gradients!\n"
|
|
"\n"
|
|
"This means backpropagation didn't work.\n"
|
|
"\n"
|
|
"💡 DEBUG STEPS:\n"
|
|
"1. Check loss.backward() is called\n"
|
|
"2. Check gradients flow through activation\n"
|
|
"3. Check gradients flow through linear layer"
|
|
)
|
|
|
|
# Verify gradients are not all zeros
|
|
assert np.any(layer.weight.grad != 0), (
|
|
"Gradients are all zeros!\n"
|
|
"\n"
|
|
"This usually means:\n"
|
|
"- ReLU killed all gradients (all outputs were negative)\n"
|
|
"- A backward function returns zeros\n"
|
|
"\n"
|
|
"💡 TRY: Print intermediate values to find where gradients die"
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# SECTION 4: COMMON MISTAKES (Educational)
|
|
# Tests that catch common student errors
|
|
# =============================================================================
|
|
|
|
class TestCommonMistakes:
|
|
"""
|
|
⚠️ COMMON MISTAKE DETECTION
|
|
|
|
These tests catch mistakes students often make.
|
|
If these fail, check the hints carefully!
|
|
"""
|
|
|
|
def test_backward_with_scalar_loss(self):
|
|
"""
|
|
⚠️ COMMON MISTAKE: Calling backward() on non-scalar
|
|
|
|
backward() should be called on the loss (a scalar).
|
|
You can't backprop from a multi-element tensor directly.
|
|
"""
|
|
from tinytorch import Tensor
|
|
|
|
x = Tensor([1.0, 2.0, 3.0], requires_grad=True)
|
|
y = x * 2
|
|
|
|
# Should be able to call backward on scalar
|
|
loss = y.sum() # scalar
|
|
loss.backward() # This should work
|
|
|
|
assert x.grad is not None, "backward() on scalar loss should compute gradients"
|
|
|
|
def test_gradient_accumulation(self):
|
|
"""
|
|
⚠️ COMMON MISTAKE: Forgetting that gradients accumulate
|
|
|
|
📖 CONCEPT: Each backward() ADDS to .grad, doesn't replace it.
|
|
This is intentional (for batch accumulation).
|
|
But you need to zero gradients between training steps!
|
|
"""
|
|
from tinytorch import Tensor
|
|
|
|
x = Tensor([1.0], requires_grad=True)
|
|
|
|
# First backward
|
|
y1 = x * 2
|
|
y1.backward()
|
|
grad1 = x.grad.copy() if hasattr(x.grad, 'copy') else np.array(x.grad)
|
|
|
|
# Second backward (gradients should accumulate)
|
|
y2 = x * 2
|
|
y2.backward()
|
|
grad2 = x.grad
|
|
|
|
# Second gradient should be double the first
|
|
assert np.allclose(grad2, grad1 * 2), (
|
|
"Gradients not accumulating!\n"
|
|
"\n"
|
|
"📖 IMPORTANT: backward() should ADD to .grad, not replace.\n"
|
|
"This enables gradient accumulation across mini-batches.\n"
|
|
"\n"
|
|
"💡 In your backward functions, use:\n"
|
|
" if tensor.grad is None:\n"
|
|
" tensor.grad = gradient\n"
|
|
" else:\n"
|
|
" tensor.grad = tensor.grad + gradient"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 70)
|
|
print("Module 05: Autograd - Progressive Tests")
|
|
print("=" * 70)
|
|
print()
|
|
print("To run these tests:")
|
|
print(" pytest tests/progressive/test_module_05_autograd.py -v")
|
|
print()
|
|
print("Or via tito:")
|
|
print(" tito module test 05")
|
|
print()
|
|
pytest.main([__file__, "-v"])
|
|
|