mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-04-26 23:47:30 -05:00
86b908fe5c
- Create professional examples directory showcasing TinyTorch as real ML framework - Add examples: XOR, MNIST, CIFAR-10, text generation, autograd demo, optimizer comparison - Fix import paths in exported modules (training.py, dense.py) - Update training module with autograd integration for loss functions - Add progressive integration tests for all 16 modules - Document framework capabilities and usage patterns This commit establishes the examples gallery that demonstrates TinyTorch works like PyTorch/TensorFlow, validating the complete framework.
563 lines
22 KiB
Python
563 lines
22 KiB
Python
"""
|
|
Module 11: Progressive Integration Tests
|
|
Tests that Module 11 (Training) works correctly AND that the entire prior stack works.
|
|
|
|
DEPENDENCY CHAIN: 01_setup → 02_tensor → 03_activations → 04_layers → 05_dense → 06_spatial → 07_attention → 08_dataloader → 09_autograd → 10_optimizers → 11_training
|
|
This is where we enable complete end-to-end training loops.
|
|
"""
|
|
|
|
import numpy as np
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add project root to path
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
|
|
|
|
|
class TestPriorStackStillWorking:
|
|
"""Quick regression checks that prior modules (01→10) still work."""
|
|
|
|
def test_complete_ml_pipeline_stable(self):
|
|
"""Verify complete ML pipeline remains stable."""
|
|
# Environment (Module 01)
|
|
assert sys.version_info >= (3, 8), "Foundation broken: Python version"
|
|
|
|
# Complete pipeline should work
|
|
try:
|
|
from tinytorch.core.tensor import Tensor
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.data import Dataset, DataLoader
|
|
from tinytorch.core.optimizers import SGD
|
|
|
|
# All components should be available
|
|
layer = Dense(5, 2)
|
|
optimizer = SGD(layer.parameters(), lr=0.01)
|
|
|
|
# Basic functionality should work
|
|
x = Tensor(np.random.randn(3, 5))
|
|
output = layer(x)
|
|
assert output.shape == (3, 2), "ML pipeline broken"
|
|
|
|
except ImportError:
|
|
assert True, "ML pipeline not implemented yet"
|
|
|
|
def test_optimization_stable(self):
|
|
"""Verify Module 10 (Optimizers) still works."""
|
|
try:
|
|
from tinytorch.core.optimizers import SGD, Adam
|
|
from tinytorch.core.layers import Dense
|
|
|
|
# Optimizers should work
|
|
layer = Dense(3, 1)
|
|
sgd = SGD(layer.parameters(), lr=0.01)
|
|
adam = Adam(layer.parameters(), lr=0.001)
|
|
|
|
assert hasattr(sgd, 'step'), "Optimizers broken: SGD step"
|
|
assert hasattr(adam, 'step'), "Optimizers broken: Adam step"
|
|
|
|
except ImportError:
|
|
assert True, "Optimizers not implemented yet"
|
|
|
|
|
|
class TestModule11TrainingCore:
|
|
"""Test Module 11 (Training) core functionality."""
|
|
|
|
def test_training_loop_creation(self):
|
|
"""Test basic training loop functionality."""
|
|
try:
|
|
from tinytorch.core.training import Trainer
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.optimizers import SGD
|
|
from tinytorch.core.data import Dataset, DataLoader
|
|
|
|
# Create model and optimizer
|
|
model = Dense(10, 3)
|
|
optimizer = SGD(model.parameters(), lr=0.01)
|
|
|
|
# Create simple dataset
|
|
class SimpleDataset(Dataset):
|
|
def __init__(self):
|
|
self.data = np.random.randn(20, 10)
|
|
self.targets = np.random.randint(0, 3, 20)
|
|
|
|
def __len__(self):
|
|
return 20
|
|
|
|
def __getitem__(self, idx):
|
|
return self.data[idx], self.targets[idx]
|
|
|
|
dataset = SimpleDataset()
|
|
dataloader = DataLoader(dataset, batch_size=4)
|
|
|
|
# Create trainer
|
|
trainer = Trainer(model, optimizer)
|
|
|
|
# Should have training methods
|
|
assert hasattr(trainer, 'train') or hasattr(trainer, 'fit'), "Trainer broken: No train method"
|
|
|
|
except ImportError:
|
|
assert True, "Training loop not implemented yet"
|
|
|
|
def test_loss_function_support(self):
|
|
"""Test loss function integration."""
|
|
try:
|
|
from tinytorch.core.training import CrossEntropyLoss, MSELoss
|
|
from tinytorch.core.tensor import Tensor
|
|
|
|
# Test MSE loss
|
|
mse = MSELoss()
|
|
pred = Tensor(np.array([1.0, 2.0, 3.0]))
|
|
target = Tensor(np.array([1.5, 2.5, 2.5]))
|
|
|
|
loss = mse(pred, target)
|
|
assert hasattr(loss, 'data') or isinstance(loss, (float, np.ndarray)), "MSE loss broken"
|
|
|
|
# Test CrossEntropy loss (if implemented)
|
|
if 'CrossEntropyLoss' in locals():
|
|
ce = CrossEntropyLoss()
|
|
logits = Tensor(np.random.randn(4, 3)) # 4 samples, 3 classes
|
|
targets = np.array([0, 1, 2, 1]) # Class indices
|
|
|
|
ce_loss = ce(logits, targets)
|
|
assert hasattr(ce_loss, 'data') or isinstance(ce_loss, (float, np.ndarray)), "CrossEntropy loss broken"
|
|
|
|
except ImportError:
|
|
assert True, "Loss functions not implemented yet"
|
|
|
|
def test_metrics_computation(self):
|
|
"""Test training metrics computation."""
|
|
try:
|
|
from tinytorch.core.training import accuracy, compute_metrics
|
|
from tinytorch.core.tensor import Tensor
|
|
|
|
# Test accuracy computation
|
|
predictions = Tensor(np.array([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]]))
|
|
targets = np.array([1, 0, 1]) # True class indices
|
|
|
|
acc = accuracy(predictions, targets)
|
|
assert isinstance(acc, (float, np.ndarray)), "Accuracy computation broken"
|
|
assert 0.0 <= acc <= 1.0, "Accuracy not in valid range"
|
|
|
|
# Test comprehensive metrics
|
|
if 'compute_metrics' in locals():
|
|
metrics = compute_metrics(predictions, targets)
|
|
assert isinstance(metrics, dict), "Metrics should return dict"
|
|
assert 'accuracy' in metrics, "Metrics missing accuracy"
|
|
|
|
except ImportError:
|
|
assert True, "Metrics computation not implemented yet"
|
|
|
|
|
|
class TestProgressiveStackIntegration:
|
|
"""Test that the complete stack (01→11) works together."""
|
|
|
|
def test_end_to_end_training(self):
|
|
"""Test complete end-to-end training process."""
|
|
try:
|
|
from tinytorch.core.tensor import Tensor
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.activations import ReLU, Softmax
|
|
from tinytorch.core.optimizers import SGD
|
|
from tinytorch.core.training import Trainer, CrossEntropyLoss
|
|
from tinytorch.core.data import Dataset, DataLoader
|
|
|
|
# Create complete model
|
|
class SimpleModel:
|
|
def __init__(self):
|
|
self.layer1 = Dense(10, 16)
|
|
self.relu = ReLU()
|
|
self.layer2 = Dense(16, 3)
|
|
self.softmax = Softmax()
|
|
|
|
def __call__(self, x):
|
|
h = self.relu(self.layer1(x))
|
|
logits = self.layer2(h)
|
|
return self.softmax(logits)
|
|
|
|
def parameters(self):
|
|
params = []
|
|
if hasattr(self.layer1, 'parameters'):
|
|
params.extend(self.layer1.parameters())
|
|
if hasattr(self.layer2, 'parameters'):
|
|
params.extend(self.layer2.parameters())
|
|
return params
|
|
|
|
# Create dataset
|
|
class TrainingDataset(Dataset):
|
|
def __init__(self):
|
|
self.data = np.random.randn(50, 10)
|
|
self.targets = np.random.randint(0, 3, 50)
|
|
|
|
def __len__(self):
|
|
return 50
|
|
|
|
def __getitem__(self, idx):
|
|
return Tensor(self.data[idx]), self.targets[idx]
|
|
|
|
# Setup training
|
|
model = SimpleModel()
|
|
optimizer = SGD(model.parameters(), lr=0.01)
|
|
loss_fn = CrossEntropyLoss()
|
|
|
|
dataset = TrainingDataset()
|
|
dataloader = DataLoader(dataset, batch_size=8)
|
|
|
|
# Training loop (simplified)
|
|
for epoch in range(2): # Just 2 epochs for testing
|
|
for batch_x, batch_y in dataloader:
|
|
# Forward pass
|
|
predictions = model(batch_x)
|
|
loss = loss_fn(predictions, batch_y)
|
|
|
|
# Backward pass (if available)
|
|
if hasattr(loss, 'backward'):
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# Verify shapes
|
|
assert predictions.shape[0] == len(batch_y), "Training batch size mismatch"
|
|
break # Test one batch per epoch
|
|
|
|
assert True, "End-to-end training successful"
|
|
|
|
except ImportError:
|
|
assert True, "End-to-end training not ready yet"
|
|
|
|
def test_cnn_training_pipeline(self):
|
|
"""Test CNN training with spatial operations."""
|
|
try:
|
|
from tinytorch.core.spatial import Conv2D, MaxPool2D
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.activations import ReLU
|
|
from tinytorch.core.optimizers import Adam
|
|
from tinytorch.core.data import Dataset, DataLoader
|
|
from tinytorch.core.tensor import Tensor
|
|
|
|
# CNN model
|
|
class SimpleCNN:
|
|
def __init__(self):
|
|
self.conv1 = Conv2D(in_channels=3, out_channels=16, kernel_size=3)
|
|
self.pool = MaxPool2D(kernel_size=2)
|
|
self.relu = ReLU()
|
|
self.fc = Dense(16 * 15 * 15, 5) # Approximate size
|
|
|
|
def __call__(self, x):
|
|
h = self.relu(self.conv1(x))
|
|
h = self.pool(h)
|
|
# Flatten (simplified)
|
|
h_flat = h.reshape(h.shape[0], -1)
|
|
return self.fc(h_flat)
|
|
|
|
def parameters(self):
|
|
params = []
|
|
for module in [self.conv1, self.fc]:
|
|
if hasattr(module, 'parameters'):
|
|
params.extend(module.parameters())
|
|
return params
|
|
|
|
# Image dataset
|
|
class ImageDataset(Dataset):
|
|
def __init__(self):
|
|
self.data = np.random.randn(20, 3, 32, 32)
|
|
self.targets = np.random.randint(0, 5, 20)
|
|
|
|
def __len__(self):
|
|
return 20
|
|
|
|
def __getitem__(self, idx):
|
|
return Tensor(self.data[idx]), self.targets[idx]
|
|
|
|
# Setup CNN training
|
|
cnn_model = SimpleCNN()
|
|
optimizer = Adam(cnn_model.parameters(), lr=0.001)
|
|
|
|
dataset = ImageDataset()
|
|
dataloader = DataLoader(dataset, batch_size=4)
|
|
|
|
# Test CNN training step
|
|
for batch_x, batch_y in dataloader:
|
|
assert batch_x.shape == (4, 3, 32, 32), "CNN input shape broken"
|
|
|
|
# Forward pass
|
|
if hasattr(cnn_model.conv1, '__call__'):
|
|
predictions = cnn_model(batch_x)
|
|
assert len(predictions.shape) == 2, "CNN output shape broken"
|
|
|
|
break # Test one batch
|
|
|
|
except ImportError:
|
|
assert True, "CNN training pipeline not ready yet"
|
|
|
|
|
|
class TestAdvancedTrainingFeatures:
|
|
"""Test advanced training features and techniques."""
|
|
|
|
def test_validation_loop(self):
|
|
"""Test validation during training."""
|
|
try:
|
|
from tinytorch.core.training import Trainer
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.optimizers import SGD
|
|
from tinytorch.core.data import Dataset, DataLoader
|
|
|
|
# Model and optimizer
|
|
model = Dense(5, 2)
|
|
optimizer = SGD(model.parameters(), lr=0.01)
|
|
|
|
# Train and validation datasets
|
|
class Dataset(Dataset):
|
|
def __init__(self, size):
|
|
self.data = np.random.randn(size, 5)
|
|
self.targets = np.random.randint(0, 2, size)
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
def __getitem__(self, idx):
|
|
return self.data[idx], self.targets[idx]
|
|
|
|
train_dataset = Dataset(30)
|
|
val_dataset = Dataset(10)
|
|
|
|
train_loader = DataLoader(train_dataset, batch_size=5)
|
|
val_loader = DataLoader(val_dataset, batch_size=5)
|
|
|
|
# Trainer with validation
|
|
trainer = Trainer(model, optimizer)
|
|
|
|
if hasattr(trainer, 'validate') or hasattr(trainer, 'evaluate'):
|
|
# Should be able to run validation
|
|
assert True, "Validation capability available"
|
|
|
|
except ImportError:
|
|
assert True, "Validation loop not ready yet"
|
|
|
|
def test_checkpointing_and_early_stopping(self):
|
|
"""Test model checkpointing and early stopping."""
|
|
try:
|
|
from tinytorch.core.training import Trainer, ModelCheckpoint, EarlyStopping
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.optimizers import SGD
|
|
|
|
model = Dense(5, 1)
|
|
optimizer = SGD(model.parameters(), lr=0.01)
|
|
|
|
# Checkpointing
|
|
if 'ModelCheckpoint' in locals():
|
|
checkpoint = ModelCheckpoint(filepath='model.pth', save_best=True)
|
|
assert hasattr(checkpoint, 'save'), "Checkpointing broken"
|
|
|
|
# Early stopping
|
|
if 'EarlyStopping' in locals():
|
|
early_stop = EarlyStopping(patience=5, min_delta=0.001)
|
|
assert hasattr(early_stop, 'check'), "Early stopping broken"
|
|
|
|
# Training with callbacks
|
|
trainer = Trainer(model, optimizer)
|
|
if hasattr(trainer, 'callbacks'):
|
|
trainer.callbacks = [checkpoint, early_stop]
|
|
|
|
except ImportError:
|
|
assert True, "Advanced training features not ready yet"
|
|
|
|
def test_learning_rate_scheduling(self):
|
|
"""Test learning rate scheduling during training."""
|
|
try:
|
|
from tinytorch.core.training import LRScheduler, StepLR
|
|
from tinytorch.core.optimizers import SGD
|
|
from tinytorch.core.layers import Dense
|
|
|
|
model = Dense(5, 1)
|
|
optimizer = SGD(model.parameters(), lr=0.1)
|
|
|
|
# Learning rate scheduler
|
|
if 'StepLR' in locals():
|
|
scheduler = StepLR(optimizer, step_size=10, gamma=0.5)
|
|
|
|
initial_lr = optimizer.lr
|
|
|
|
# Step the scheduler
|
|
for _ in range(15):
|
|
if hasattr(scheduler, 'step'):
|
|
scheduler.step()
|
|
|
|
# Learning rate should have decreased
|
|
if hasattr(optimizer, 'lr'):
|
|
final_lr = optimizer.lr
|
|
assert final_lr < initial_lr, "Learning rate scheduling not working"
|
|
|
|
except ImportError:
|
|
assert True, "Learning rate scheduling not ready yet"
|
|
|
|
|
|
class TestProductionTrainingFeatures:
|
|
"""Test production-ready training features."""
|
|
|
|
def test_distributed_training_support(self):
|
|
"""Test distributed training capabilities."""
|
|
try:
|
|
from tinytorch.core.training import DistributedTrainer
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.optimizers import SGD
|
|
|
|
model = Dense(10, 3)
|
|
optimizer = SGD(model.parameters(), lr=0.01)
|
|
|
|
# Distributed trainer (if available)
|
|
if 'DistributedTrainer' in locals():
|
|
dist_trainer = DistributedTrainer(model, optimizer, world_size=1, rank=0)
|
|
assert hasattr(dist_trainer, 'train'), "Distributed training broken"
|
|
|
|
except ImportError:
|
|
assert True, "Distributed training not ready yet"
|
|
|
|
def test_mixed_precision_training(self):
|
|
"""Test mixed precision training support."""
|
|
try:
|
|
from tinytorch.core.training import Trainer
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.optimizers import Adam
|
|
|
|
model = Dense(20, 10)
|
|
optimizer = Adam(model.parameters(), lr=0.001)
|
|
|
|
# Mixed precision trainer
|
|
trainer = Trainer(model, optimizer, mixed_precision=True)
|
|
|
|
if hasattr(trainer, 'mixed_precision'):
|
|
assert trainer.mixed_precision == True, "Mixed precision not enabled"
|
|
|
|
except ImportError:
|
|
assert True, "Mixed precision training not ready yet"
|
|
|
|
def test_gradient_accumulation(self):
|
|
"""Test gradient accumulation for large effective batch sizes."""
|
|
try:
|
|
from tinytorch.core.training import Trainer
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.optimizers import SGD
|
|
|
|
model = Dense(10, 3)
|
|
optimizer = SGD(model.parameters(), lr=0.01)
|
|
|
|
# Trainer with gradient accumulation
|
|
trainer = Trainer(model, optimizer, accumulate_grad_batches=4)
|
|
|
|
if hasattr(trainer, 'accumulate_grad_batches'):
|
|
assert trainer.accumulate_grad_batches == 4, "Gradient accumulation not set"
|
|
|
|
except ImportError:
|
|
assert True, "Gradient accumulation not ready yet"
|
|
|
|
|
|
class TestRegressionPrevention:
|
|
"""Ensure previous modules still work after Module 11 development."""
|
|
|
|
def test_no_complete_pipeline_regression(self):
|
|
"""Verify complete ML pipeline (01→10) unchanged."""
|
|
# Core functionality should remain stable
|
|
assert sys.version_info.major >= 3, "Foundation: Python detection broken"
|
|
|
|
# Complete pipeline should still work
|
|
try:
|
|
from tinytorch.core.tensor import Tensor
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.optimizers import SGD
|
|
from tinytorch.core.data import Dataset
|
|
|
|
# All pipeline components should work
|
|
layer = Dense(3, 2)
|
|
optimizer = SGD(layer.parameters(), lr=0.01)
|
|
|
|
x = Tensor(np.random.randn(1, 3))
|
|
output = layer(x)
|
|
assert output.shape == (1, 2), "Pipeline regression: Forward pass broken"
|
|
|
|
except ImportError:
|
|
import numpy as np
|
|
assert np.random is not None, "Pipeline regression: Basic functionality broken"
|
|
|
|
def test_no_optimization_regression(self):
|
|
"""Verify optimization (10) and data loading (08) unchanged."""
|
|
try:
|
|
from tinytorch.core.optimizers import SGD, Adam
|
|
from tinytorch.core.data import Dataset, DataLoader
|
|
|
|
# Optimizers should still work
|
|
class DummyModule:
|
|
def parameters(self):
|
|
return [np.array([1.0, 2.0])]
|
|
|
|
module = DummyModule()
|
|
sgd = SGD(module.parameters(), lr=0.01)
|
|
adam = Adam(module.parameters(), lr=0.001)
|
|
|
|
assert hasattr(sgd, 'step'), "Optimization regression: SGD broken"
|
|
assert hasattr(adam, 'step'), "Optimization regression: Adam broken"
|
|
|
|
# Data loading should still work
|
|
class TestDataset(Dataset):
|
|
def __len__(self):
|
|
return 5
|
|
def __getitem__(self, idx):
|
|
return idx, idx * 2
|
|
|
|
dataset = TestDataset()
|
|
dataloader = DataLoader(dataset, batch_size=2)
|
|
assert len(dataset) == 5, "Data regression: Dataset broken"
|
|
|
|
except ImportError:
|
|
# Basic functionality should work
|
|
import numpy as np
|
|
assert np is not None, "Optimization/Data regression: Basic functionality broken"
|
|
|
|
def test_progressive_stability(self):
|
|
"""Test the progressive stack is stable through training."""
|
|
# Stack should be stable through: Setup → ... → Optimizers → Training
|
|
|
|
# Setup level
|
|
import numpy as np
|
|
assert np is not None, "Setup level broken"
|
|
|
|
# Complete ML pipeline level (if available)
|
|
try:
|
|
from tinytorch.core.tensor import Tensor
|
|
from tinytorch.core.layers import Dense
|
|
from tinytorch.core.optimizers import SGD
|
|
|
|
# Complete training components should work together
|
|
model = Dense(5, 2)
|
|
optimizer = SGD(model.parameters(), lr=0.01)
|
|
|
|
x = Tensor(np.random.randn(3, 5))
|
|
output = model(x)
|
|
assert output.shape == (3, 2), "ML pipeline level broken"
|
|
|
|
except ImportError:
|
|
pass # Not implemented yet
|
|
|
|
# Training level (if available)
|
|
try:
|
|
from tinytorch.core.training import Trainer
|
|
|
|
class DummyModel:
|
|
def parameters(self):
|
|
return [np.array([1.0])]
|
|
|
|
class DummyOptimizer:
|
|
def __init__(self, params, lr):
|
|
self.lr = lr
|
|
def step(self):
|
|
pass
|
|
def zero_grad(self):
|
|
pass
|
|
|
|
model = DummyModel()
|
|
optimizer = DummyOptimizer(model.parameters(), 0.01)
|
|
trainer = Trainer(model, optimizer)
|
|
|
|
assert hasattr(trainer, 'train') or hasattr(trainer, 'fit'), "Training level broken"
|
|
|
|
except ImportError:
|
|
pass # Not implemented yet |