mirror of
https://github.com/harvard-edge/cs249r_book.git
synced 2026-05-03 00:07:08 -05:00
test_autograd_integration() and test_loss_backward_integration() now gracefully skip if requires_grad is not available (i.e., autograd hasn't been enabled yet). This prevents false failures when running integration tests before Module 06 has been completed.
349 lines
12 KiB
Python
349 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Module Dependency Integration Testing
|
|
Tests how each module interfaces with modules that came before it
|
|
"""
|
|
|
|
# Module dependency graph for TinyTorch
|
|
# Current module structure:
|
|
# 01_tensor, 02_activations, 03_layers, 04_losses, 05_dataloader,
|
|
# 06_autograd, 07_optimizers, 08_training, 09_convolutions,
|
|
# 10_tokenization, 11_embeddings, 12_attention, 13_transformers,
|
|
# 14_profiling, 15_quantization, 16_compression, 17_acceleration,
|
|
# 18_memoization, 19_benchmarking, 20_capstone
|
|
MODULE_DEPENDENCIES = {
|
|
"01_tensor": [], # No dependencies - foundation
|
|
"02_activations": ["01_tensor"], # Needs Tensor
|
|
"03_layers": ["01_tensor"], # Needs Tensor
|
|
"04_losses": ["01_tensor"], # Needs Tensor
|
|
"05_dataloader": ["01_tensor"], # Needs Tensor
|
|
"06_autograd": ["01_tensor"], # Core dependency on Tensor
|
|
"07_optimizers": ["01_tensor", "06_autograd"], # Needs Tensor and autograd
|
|
"08_training": ["01_tensor", "05_dataloader", "06_autograd", "07_optimizers"], # Training loop deps
|
|
"09_convolutions": ["01_tensor", "03_layers"], # Needs Tensor and Layer base
|
|
"10_tokenization": ["01_tensor"], # Needs Tensor
|
|
"11_embeddings": ["01_tensor"], # Needs Tensor
|
|
"12_attention": ["01_tensor", "03_layers"], # Needs Tensor, Layer
|
|
"13_transformers": ["01_tensor", "03_layers", "12_attention"], # Full attention stack
|
|
"14_profiling": ["01_tensor"], # Performance analysis
|
|
"15_quantization": ["01_tensor"], # Optimization techniques
|
|
"16_compression": ["01_tensor"], # Optimization techniques
|
|
"17_acceleration": ["01_tensor"], # Runtime optimization (general)
|
|
"18_memoization": ["01_tensor"], # Runtime optimization (transformer-specific)
|
|
"19_benchmarking": ["01_tensor"], # Performance testing
|
|
"20_capstone": ["01_tensor", "09_convolutions", "13_transformers"] # Full stack
|
|
}
|
|
|
|
def get_module_integration_tests(module_name: str):
|
|
"""
|
|
Get integration tests based on module dependencies.
|
|
Returns a list of test functions to run.
|
|
"""
|
|
tests = []
|
|
|
|
# Get dependencies for this module
|
|
deps = MODULE_DEPENDENCIES.get(module_name, [])
|
|
|
|
# Generate tests based on dependencies
|
|
if "02_tensor" in deps:
|
|
tests.append(("test_tensor_integration", test_tensor_integration))
|
|
|
|
if "04_layers" in deps:
|
|
tests.append(("test_layer_integration", test_layer_integration))
|
|
|
|
if "05_dataloader" in deps:
|
|
tests.append(("test_dataloader_integration", test_dataloader_integration))
|
|
|
|
if "06_autograd" in deps:
|
|
tests.append(("test_autograd_integration", test_autograd_integration))
|
|
|
|
if "07_optimizers" in deps:
|
|
tests.append(("test_optimizer_integration", test_optimizer_integration))
|
|
|
|
# Module-specific integration tests
|
|
if module_name == "05_dataloader":
|
|
tests.append(("test_dataloader_with_tensor", test_dataloader_with_tensor))
|
|
tests.append(("test_dataloader_with_batching", test_dataloader_with_batching))
|
|
tests.append(("test_dataloader_pipeline", test_dataloader_pipeline))
|
|
|
|
elif module_name == "09_convolutions":
|
|
tests.append(("test_conv2d_with_tensor", test_conv2d_with_tensor))
|
|
tests.append(("test_pooling_integration", test_pooling_integration))
|
|
|
|
elif module_name == "07_attention":
|
|
tests.append(("test_attention_with_dense", test_attention_with_dense))
|
|
tests.append(("test_multihead_integration", test_multihead_integration))
|
|
|
|
elif module_name == "12_training":
|
|
tests.append(("test_training_loop_integration", test_training_loop_integration))
|
|
tests.append(("test_loss_backward_integration", test_loss_backward_integration))
|
|
|
|
return tests
|
|
|
|
|
|
# Base integration tests that check module interfaces
|
|
def test_tensor_integration():
|
|
"""Test that Tensor works as expected for dependent modules."""
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
# Test tensor creation
|
|
t = Tensor(np.array([1, 2, 3]))
|
|
assert t.shape == (3,), "Tensor shape should work"
|
|
assert t.data is not None, "Tensor should have data"
|
|
|
|
# Test tensor operations needed by other modules
|
|
t2 = Tensor(np.array([4, 5, 6]))
|
|
result = t.data + t2.data # Many modules need element-wise ops
|
|
assert result.shape == (3,), "Element-wise ops should preserve shape"
|
|
|
|
|
|
def test_layer_integration():
|
|
"""Test Layer base class interface."""
|
|
from tinytorch.core.layers import Layer
|
|
|
|
# Test that Layer exists and has expected interface
|
|
assert hasattr(Layer, 'forward'), "Layer should have forward method"
|
|
assert hasattr(Layer, '__call__'), "Layer should be callable"
|
|
|
|
# Test basic layer creation
|
|
layer = Layer()
|
|
assert layer is not None, "Should create Layer instance"
|
|
|
|
|
|
def test_dense_integration():
|
|
"""Test Dense layer integration with Tensor."""
|
|
from tinytorch.core.layers import Linear
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
# Test Dense with Tensor input
|
|
layer = Linear(10, 5)
|
|
x = Tensor(np.random.randn(32, 10))
|
|
output = layer(x)
|
|
|
|
assert output.shape == (32, 5), "Dense should produce correct shape"
|
|
assert isinstance(output, Tensor), "Dense should return Tensor"
|
|
|
|
|
|
def test_dense_with_tensor():
|
|
"""Test that Dense properly uses Tensor for weights/bias."""
|
|
from tinytorch.core.layers import Linear
|
|
from tinytorch.core.tensor import Tensor
|
|
|
|
layer = Linear(10, 5)
|
|
|
|
# Check weights are Tensors
|
|
assert isinstance(layer.weight, Tensor), "Weights should be Tensor"
|
|
assert layer.weight.shape == (10, 5), "Weight shape should match layer dims"
|
|
# Bias may or may not exist depending on implementation
|
|
if hasattr(layer, 'bias') and layer.bias is not None:
|
|
assert isinstance(layer.bias, Tensor), "Bias should be Tensor"
|
|
|
|
|
|
def test_dense_with_activations():
|
|
"""Test Dense layer works with activation functions."""
|
|
from tinytorch.core.layers import Linear
|
|
from tinytorch.core.activations import ReLU, Sigmoid
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
# Build small network: Dense -> ReLU -> Dense -> Sigmoid
|
|
layer1 = Linear(10, 20)
|
|
relu = ReLU()
|
|
layer2 = Linear(20, 1)
|
|
sigmoid = Sigmoid()
|
|
|
|
# Forward pass
|
|
x = Tensor(np.random.randn(16, 10))
|
|
h1 = layer1(x)
|
|
h1_activated = relu(h1)
|
|
output = layer2(h1_activated)
|
|
final = sigmoid(output)
|
|
|
|
# Check shapes preserved through network
|
|
assert h1.shape == (16, 20), "First layer output shape"
|
|
assert h1_activated.shape == (16, 20), "ReLU preserves shape"
|
|
assert output.shape == (16, 1), "Second layer output shape"
|
|
assert final.shape == (16, 1), "Sigmoid preserves shape"
|
|
|
|
# Check sigmoid output range
|
|
assert np.all(final.data >= 0) and np.all(final.data <= 1), "Sigmoid outputs in [0,1]"
|
|
|
|
|
|
def test_multi_layer_network():
|
|
"""Test building multi-layer networks with Dense."""
|
|
from tinytorch.core.layers import Linear
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
# Build 3-layer network
|
|
layers = [
|
|
Linear(784, 128),
|
|
Linear(128, 64),
|
|
Linear(64, 10)
|
|
]
|
|
|
|
# Forward pass through all layers
|
|
x = Tensor(np.random.randn(32, 784))
|
|
|
|
for i, layer in enumerate(layers):
|
|
x = layer(x)
|
|
if i == 0:
|
|
assert x.shape == (32, 128), f"Layer {i} shape"
|
|
elif i == 1:
|
|
assert x.shape == (32, 64), f"Layer {i} shape"
|
|
elif i == 2:
|
|
assert x.shape == (32, 10), f"Layer {i} shape"
|
|
|
|
assert x.shape == (32, 10), "Final output shape should be (32, 10)"
|
|
|
|
|
|
def test_conv2d_with_tensor():
|
|
"""Test Conv2d integration with Tensor."""
|
|
from tinytorch.core.spatial import Conv2d
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
# Create Conv2d layer
|
|
conv = Conv2d(in_channels=3, out_channels=16, kernel_size=3)
|
|
|
|
# Test with image tensor (batch, channels, height, width)
|
|
x = Tensor(np.random.randn(8, 3, 32, 32))
|
|
output = conv(x)
|
|
|
|
# Check output shape (with valid padding, output is smaller)
|
|
assert output.shape[0] == 8, "Batch size preserved"
|
|
assert output.shape[1] == 16, "Output channels correct"
|
|
|
|
|
|
def test_pooling_integration():
|
|
"""Test pooling layers work with Conv2d output."""
|
|
from tinytorch.core.spatial import Conv2d, MaxPool2d
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
conv = Conv2d(3, 32, kernel_size=3, padding=1)
|
|
pool = MaxPool2d(kernel_size=2, stride=2)
|
|
|
|
x = Tensor(np.random.randn(4, 3, 28, 28))
|
|
conv_out = conv(x)
|
|
pool_out = pool(conv_out)
|
|
|
|
# Pooling should reduce spatial dimensions by half
|
|
assert pool_out.shape[2] == conv_out.shape[2] // 2
|
|
assert pool_out.shape[3] == conv_out.shape[3] // 2
|
|
|
|
|
|
def test_attention_with_dense():
|
|
"""Test attention mechanism uses Dense layers."""
|
|
from tinytorch.core.attention import MultiHeadAttention
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
attention = MultiHeadAttention(embed_dim=64, num_heads=4)
|
|
x = Tensor(np.random.randn(2, 10, 64)) # (batch, seq_len, embed_dim)
|
|
|
|
output = attention(x)
|
|
assert output.shape == x.shape, "Attention preserves shape"
|
|
|
|
|
|
def test_multihead_integration():
|
|
"""Test multi-head attention integration."""
|
|
from tinytorch.core.attention import MultiHeadAttention
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
mha = MultiHeadAttention(embed_dim=64, num_heads=8)
|
|
x = Tensor(np.random.randn(2, 10, 64))
|
|
|
|
output = mha(x)
|
|
assert output.shape == x.shape, "MHA preserves input shape"
|
|
|
|
|
|
def test_autograd_integration():
|
|
"""Test autograd system with Tensor.
|
|
|
|
NOTE: This test requires autograd to be enabled (Module 06+).
|
|
It will skip if requires_grad is not available.
|
|
"""
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
# Check if autograd is enabled (requires_grad parameter available)
|
|
try:
|
|
x = Tensor(np.array([[1, 2], [3, 4]]), requires_grad=True)
|
|
except TypeError:
|
|
# requires_grad not available - autograd not enabled yet
|
|
return # Skip test
|
|
|
|
assert hasattr(x, 'grad'), "Tensor should have grad attribute"
|
|
assert x.requires_grad == True, "Should track gradients"
|
|
|
|
|
|
def test_optimizer_integration():
|
|
"""Test optimizers work with layers."""
|
|
from tinytorch.core.optimizers import SGD
|
|
from tinytorch.core.layers import Linear
|
|
|
|
layer = Linear(10, 5)
|
|
params = layer.parameters()
|
|
optimizer = SGD(params, lr=0.01)
|
|
|
|
# Test optimizer has params
|
|
assert len(params) > 0, "Layer should have parameters"
|
|
|
|
|
|
def test_training_loop_integration():
|
|
"""Test training loop integrates optimizer and autograd."""
|
|
from tinytorch.core.layers import Linear
|
|
from tinytorch.core.optimizers import SGD
|
|
from tinytorch.core.losses import MSELoss
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
# Simple model
|
|
model = Linear(10, 1)
|
|
params = model.parameters()
|
|
optimizer = SGD(params, lr=0.01)
|
|
loss_fn = MSELoss()
|
|
|
|
# Dummy data
|
|
X = Tensor(np.random.randn(32, 10))
|
|
y = Tensor(np.random.randn(32, 1))
|
|
|
|
# One training step
|
|
predictions = model(X)
|
|
loss = loss_fn(predictions, y)
|
|
|
|
# Loss should be computed
|
|
assert loss is not None, "Loss should be computed"
|
|
|
|
|
|
def test_loss_backward_integration():
|
|
"""Test loss functions integrate with autograd.
|
|
|
|
NOTE: This test requires autograd to be enabled (Module 06+).
|
|
It will skip if requires_grad is not available.
|
|
"""
|
|
from tinytorch.core.losses import MSELoss
|
|
from tinytorch.core.tensor import Tensor
|
|
import numpy as np
|
|
|
|
loss_fn = MSELoss()
|
|
|
|
# Check if autograd is enabled (requires_grad parameter available)
|
|
try:
|
|
predictions = Tensor(np.array([1.0, 2.0, 3.0]), requires_grad=True)
|
|
except TypeError:
|
|
# requires_grad not available - autograd not enabled yet
|
|
return # Skip test
|
|
|
|
targets = Tensor(np.array([1.5, 2.5, 3.5]))
|
|
|
|
loss = loss_fn(predictions, targets)
|
|
|
|
# Test backward pass
|
|
if hasattr(loss, 'backward'):
|
|
loss.backward()
|