mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-04-29 23:40:02 -05:00
Add MSEBackward and organize comprehensive test suite
New Features: - Add MSEBackward gradient computation for regression tasks - Patch MSELoss in enable_autograd() for gradient tracking - All 3 loss functions now support autograd: MSE, BCE, CrossEntropy Test Suite Organization: - Reorganize tests/ into focused directories - Create tests/integration/ for cross-module tests - Create tests/05_autograd/ for autograd edge cases - Create tests/debugging/ for common student pitfalls - Add comprehensive tests/README.md explaining test philosophy Integration Tests: - Move test_gradient_flow.py to integration/ - 20 comprehensive gradient flow tests - Tests cover: tensors, layers, activations, losses, optimizers - Tests validate: basic ops, chain rule, broadcasting, training loops - 19/20 tests passing (MSE now fixed!) Results: ✅ Perceptron learns: 50% → 93% accuracy ✅ Clean test organization guides future development ✅ Tests catch the exact bugs that broke training Pedagogical Value: - Test organization teaches testing best practices - Gradient flow tests show what integration testing catches - Sets foundation for debugging/diagnostic tests
This commit is contained in:
51
tinytorch/core/autograd.py
generated
51
tinytorch/core/autograd.py
generated
@@ -15,8 +15,8 @@
|
||||
# ║ happens! The tinytorch/ directory is just the compiled output. ║
|
||||
# ╚═══════════════════════════════════════════════════════════════════════════════╝
|
||||
# %% auto 0
|
||||
__all__ = ['Function', 'AddBackward', 'MulBackward', 'MatmulBackward', 'SumBackward', 'SigmoidBackward', 'BCEBackward',
|
||||
'enable_autograd']
|
||||
__all__ = ['Function', 'AddBackward', 'MulBackward', 'MatmulBackward', 'SumBackward', 'SigmoidBackward', 'MSEBackward',
|
||||
'BCEBackward', 'enable_autograd']
|
||||
|
||||
# %% ../../modules/source/05_autograd/autograd_dev.ipynb 1
|
||||
import numpy as np
|
||||
@@ -271,6 +271,32 @@ class SigmoidBackward(Function):
|
||||
return None,
|
||||
|
||||
# %% ../../modules/source/05_autograd/autograd_dev.ipynb 21
|
||||
class MSEBackward(Function):
|
||||
"""
|
||||
Gradient computation for Mean Squared Error Loss.
|
||||
|
||||
MSE: L = mean((predictions - targets)²)
|
||||
Derivative: ∂L/∂predictions = 2 * (predictions - targets) / N
|
||||
"""
|
||||
|
||||
def __init__(self, predictions, targets):
|
||||
"""Initialize with predictions and targets."""
|
||||
super().__init__(predictions)
|
||||
self.targets_data = targets.data
|
||||
self.num_samples = np.size(targets.data)
|
||||
|
||||
def apply(self, grad_output):
|
||||
"""Compute gradient for MSE loss."""
|
||||
predictions, = self.saved_tensors
|
||||
|
||||
if isinstance(predictions, Tensor) and predictions.requires_grad:
|
||||
# Gradient: 2 * (predictions - targets) / N
|
||||
grad = 2.0 * (predictions.data - self.targets_data) / self.num_samples
|
||||
|
||||
return grad * grad_output,
|
||||
return None,
|
||||
|
||||
# %% ../../modules/source/05_autograd/autograd_dev.ipynb 22
|
||||
class BCEBackward(Function):
|
||||
"""
|
||||
Gradient computation for Binary Cross-Entropy Loss.
|
||||
@@ -300,7 +326,7 @@ class BCEBackward(Function):
|
||||
return grad * grad_output,
|
||||
return None,
|
||||
|
||||
# %% ../../modules/source/05_autograd/autograd_dev.ipynb 22
|
||||
# %% ../../modules/source/05_autograd/autograd_dev.ipynb 23
|
||||
def enable_autograd():
|
||||
"""
|
||||
Enable gradient tracking for all Tensor operations.
|
||||
@@ -502,11 +528,12 @@ def enable_autograd():
|
||||
# Patch activations and losses to track gradients
|
||||
try:
|
||||
from tinytorch.core.activations import Sigmoid
|
||||
from tinytorch.core.losses import BinaryCrossEntropyLoss
|
||||
from tinytorch.core.losses import BinaryCrossEntropyLoss, MSELoss
|
||||
|
||||
# Store original methods
|
||||
_original_sigmoid_forward = Sigmoid.forward
|
||||
_original_bce_forward = BinaryCrossEntropyLoss.forward
|
||||
_original_mse_forward = MSELoss.forward
|
||||
|
||||
def tracked_sigmoid_forward(self, x):
|
||||
"""Sigmoid with gradient tracking."""
|
||||
@@ -537,9 +564,25 @@ def enable_autograd():
|
||||
|
||||
return result
|
||||
|
||||
def tracked_mse_forward(self, predictions, targets):
|
||||
"""MSE loss with gradient tracking."""
|
||||
# Compute MSE loss
|
||||
diff = predictions.data - targets.data
|
||||
squared_diff = diff ** 2
|
||||
mse = np.mean(squared_diff)
|
||||
|
||||
result = Tensor(mse)
|
||||
|
||||
if predictions.requires_grad:
|
||||
result.requires_grad = True
|
||||
result._grad_fn = MSEBackward(predictions, targets)
|
||||
|
||||
return result
|
||||
|
||||
# Install patched methods
|
||||
Sigmoid.forward = tracked_sigmoid_forward
|
||||
BinaryCrossEntropyLoss.forward = tracked_bce_forward
|
||||
MSELoss.forward = tracked_mse_forward
|
||||
|
||||
except ImportError:
|
||||
# Activations/losses not yet available (happens during module development)
|
||||
|
||||
Reference in New Issue
Block a user