mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-07-21 06:52:26 -05:00
Vijay Janapa Reddi
ec7168dc90
Added all module development files to modules/XX_name/ directories:
Module notebooks and scripts:
- 18 modules with .ipynb and .py files (01-20, excluding some gaps)
- Moved from modules/source/ to direct module directories
- Includes tensor, autograd, layers, transformers, optimization modules
Module README files:
- Added README.md for modules with additional documentation
- Complements ABOUT.md files added earlier
This completes the module restructuring:
- Before: modules/source/XX_name/*_dev.{py,ipynb}
- After: modules/XX_name/*_dev.{py,ipynb}
All development happens directly in numbered module directories now.
52 KiB
52 KiB
In [ ]:
#| default_exp data.loader
#| exportIn [ ]:
#| export
# Essential imports for data loading
import numpy as np
import random
from typing import Iterator, Tuple, List, Optional, Union
from abc import ABC, abstractmethod
# Import real Tensor class from tinytorch package
from tinytorch.core.tensor import TensorIn [ ]:
#| export
class Dataset(ABC):
"""
Abstract base class for all datasets.
Provides the fundamental interface that all datasets must implement:
- __len__(): Returns the total number of samples
- __getitem__(idx): Returns the sample at given index
TODO: Implement the abstract Dataset base class
APPROACH:
1. Use ABC (Abstract Base Class) to define interface
2. Mark methods as @abstractmethod to force implementation
3. Provide clear docstrings for subclasses
EXAMPLE:
>>> class MyDataset(Dataset):
... def __len__(self): return 100
... def __getitem__(self, idx): return idx
>>> dataset = MyDataset()
>>> print(len(dataset)) # 100
>>> print(dataset[42]) # 42
HINT: Abstract methods force subclasses to implement core functionality
"""
### BEGIN SOLUTION
@abstractmethod
def __len__(self) -> int:
"""
Return the total number of samples in the dataset.
This method must be implemented by all subclasses to enable
len(dataset) calls and batch size calculations.
"""
pass
@abstractmethod
def __getitem__(self, idx: int):
"""
Return the sample at the given index.
Args:
idx: Index of the sample to retrieve (0 <= idx < len(dataset))
Returns:
The sample at index idx. Format depends on the dataset implementation.
Could be (data, label) tuple, single tensor, etc.
"""
pass
### END SOLUTIONIn [ ]:
def test_unit_dataset():
"""🔬 Test Dataset abstract base class."""
print("🔬 Unit Test: Dataset Abstract Base Class...")
# Test that Dataset is properly abstract
try:
dataset = Dataset()
assert False, "Should not be able to instantiate abstract Dataset"
except TypeError:
print("✅ Dataset is properly abstract")
# Test concrete implementation
class TestDataset(Dataset):
def __init__(self, size):
self.size = size
def __len__(self):
return self.size
def __getitem__(self, idx):
return f"item_{idx}"
dataset = TestDataset(10)
assert len(dataset) == 10
assert dataset[0] == "item_0"
assert dataset[9] == "item_9"
print("✅ Dataset interface works correctly!")
if __name__ == "__main__":
test_unit_dataset()In [ ]:
#| export
class TensorDataset(Dataset):
"""
Dataset wrapping tensors for supervised learning.
Each sample is a tuple of tensors from the same index across all input tensors.
All tensors must have the same size in their first dimension.
TODO: Implement TensorDataset for tensor-based data
APPROACH:
1. Store all input tensors
2. Validate they have same first dimension (number of samples)
3. Return tuple of tensor slices for each index
EXAMPLE:
>>> features = Tensor([[1, 2], [3, 4], [5, 6]]) # 3 samples, 2 features each
>>> labels = Tensor([0, 1, 0]) # 3 labels
>>> dataset = TensorDataset(features, labels)
>>> print(len(dataset)) # 3
>>> print(dataset[1]) # (Tensor([3, 4]), Tensor(1))
HINTS:
- Use *tensors to accept variable number of tensor arguments
- Check all tensors have same length in dimension 0
- Return tuple of tensor[idx] for all tensors
"""
def __init__(self, *tensors):
"""
Create dataset from multiple tensors.
Args:
*tensors: Variable number of Tensor objects
All tensors must have the same size in their first dimension.
"""
### BEGIN SOLUTION
assert len(tensors) > 0, "Must provide at least one tensor"
# Store all tensors
self.tensors = tensors
# Validate all tensors have same first dimension
first_size = len(tensors[0].data) # Size of first dimension
for i, tensor in enumerate(tensors):
if len(tensor.data) != first_size:
raise ValueError(
f"All tensors must have same size in first dimension. "
f"Tensor 0: {first_size}, Tensor {i}: {len(tensor.data)}"
)
### END SOLUTION
def __len__(self) -> int:
"""Return number of samples (size of first dimension)."""
### BEGIN SOLUTION
return len(self.tensors[0].data)
### END SOLUTION
def __getitem__(self, idx: int) -> Tuple[Tensor, ...]:
"""
Return tuple of tensor slices at given index.
Args:
idx: Sample index
Returns:
Tuple containing tensor[idx] for each input tensor
"""
### BEGIN SOLUTION
if idx >= len(self) or idx < 0:
raise IndexError(f"Index {idx} out of range for dataset of size {len(self)}")
# Return tuple of slices from all tensors
return tuple(Tensor(tensor.data[idx]) for tensor in self.tensors)
### END SOLUTIONIn [ ]:
def test_unit_tensordataset():
"""🔬 Test TensorDataset implementation."""
print("🔬 Unit Test: TensorDataset...")
# Test basic functionality
features = Tensor([[1, 2], [3, 4], [5, 6]]) # 3 samples, 2 features
labels = Tensor([0, 1, 0]) # 3 labels
dataset = TensorDataset(features, labels)
# Test length
assert len(dataset) == 3, f"Expected length 3, got {len(dataset)}"
# Test indexing
sample = dataset[0]
assert len(sample) == 2, "Should return tuple with 2 tensors"
assert np.array_equal(sample[0].data, [1, 2]), f"Wrong features: {sample[0].data}"
assert sample[1].data == 0, f"Wrong label: {sample[1].data}"
sample = dataset[1]
assert np.array_equal(sample[1].data, 1), f"Wrong label at index 1: {sample[1].data}"
# Test error handling
try:
dataset[10] # Out of bounds
assert False, "Should raise IndexError for out of bounds access"
except IndexError:
pass
# Test mismatched tensor sizes
try:
bad_features = Tensor([[1, 2], [3, 4]]) # Only 2 samples
bad_labels = Tensor([0, 1, 0]) # 3 labels - mismatch!
TensorDataset(bad_features, bad_labels)
assert False, "Should raise error for mismatched tensor sizes"
except ValueError:
pass
print("✅ TensorDataset works correctly!")
if __name__ == "__main__":
test_unit_tensordataset()In [ ]:
#| export
class DataLoader:
"""
Data loader with batching and shuffling support.
Wraps a dataset to provide batched iteration with optional shuffling.
Essential for efficient training with mini-batch gradient descent.
TODO: Implement DataLoader with batching and shuffling
APPROACH:
1. Store dataset, batch_size, and shuffle settings
2. Create iterator that groups samples into batches
3. Handle shuffling by randomizing indices
4. Collate individual samples into batch tensors
EXAMPLE:
>>> dataset = TensorDataset(Tensor([[1,2], [3,4], [5,6]]), Tensor([0,1,0]))
>>> loader = DataLoader(dataset, batch_size=2, shuffle=True)
>>> for batch in loader:
... features_batch, labels_batch = batch
... print(f"Features: {features_batch.shape}, Labels: {labels_batch.shape}")
HINTS:
- Use random.shuffle() for index shuffling
- Group consecutive samples into batches
- Stack individual tensors using np.stack()
"""
def __init__(self, dataset: Dataset, batch_size: int, shuffle: bool = False):
"""
Create DataLoader for batched iteration.
Args:
dataset: Dataset to load from
batch_size: Number of samples per batch
shuffle: Whether to shuffle data each epoch
"""
### BEGIN SOLUTION
self.dataset = dataset
self.batch_size = batch_size
self.shuffle = shuffle
### END SOLUTION
def __len__(self) -> int:
"""Return number of batches per epoch."""
### BEGIN SOLUTION
# Calculate number of complete batches
return (len(self.dataset) + self.batch_size - 1) // self.batch_size
### END SOLUTION
def __iter__(self) -> Iterator:
"""Return iterator over batches."""
### BEGIN SOLUTION
# Create list of indices
indices = list(range(len(self.dataset)))
# Shuffle if requested
if self.shuffle:
random.shuffle(indices)
# Yield batches
for i in range(0, len(indices), self.batch_size):
batch_indices = indices[i:i + self.batch_size]
batch = [self.dataset[idx] for idx in batch_indices]
# Collate batch - convert list of tuples to tuple of tensors
yield self._collate_batch(batch)
### END SOLUTION
def _collate_batch(self, batch: List[Tuple[Tensor, ...]]) -> Tuple[Tensor, ...]:
"""
Collate individual samples into batch tensors.
Args:
batch: List of sample tuples from dataset
Returns:
Tuple of batched tensors
"""
### BEGIN SOLUTION
if len(batch) == 0:
return ()
# Determine number of tensors per sample
num_tensors = len(batch[0])
# Group tensors by position
batched_tensors = []
for tensor_idx in range(num_tensors):
# Extract all tensors at this position
tensor_list = [sample[tensor_idx].data for sample in batch]
# Stack into batch tensor
batched_data = np.stack(tensor_list, axis=0)
batched_tensors.append(Tensor(batched_data))
return tuple(batched_tensors)
### END SOLUTIONIn [ ]:
def test_unit_dataloader():
"""🔬 Test DataLoader implementation."""
print("🔬 Unit Test: DataLoader...")
# Create test dataset
features = Tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]) # 5 samples
labels = Tensor([0, 1, 0, 1, 0])
dataset = TensorDataset(features, labels)
# Test basic batching (no shuffle)
loader = DataLoader(dataset, batch_size=2, shuffle=False)
# Test length calculation
assert len(loader) == 3, f"Expected 3 batches, got {len(loader)}" # ceil(5/2) = 3
batches = list(loader)
assert len(batches) == 3, f"Expected 3 batches, got {len(batches)}"
# Test first batch
batch_features, batch_labels = batches[0]
assert batch_features.data.shape == (2, 2), f"Wrong batch features shape: {batch_features.data.shape}"
assert batch_labels.data.shape == (2,), f"Wrong batch labels shape: {batch_labels.data.shape}"
# Test last batch (should have 1 sample)
batch_features, batch_labels = batches[2]
assert batch_features.data.shape == (1, 2), f"Wrong last batch features shape: {batch_features.data.shape}"
assert batch_labels.data.shape == (1,), f"Wrong last batch labels shape: {batch_labels.data.shape}"
# Test that data is preserved
assert np.array_equal(batches[0][0].data[0], [1, 2]), "First sample should be [1,2]"
assert batches[0][1].data[0] == 0, "First label should be 0"
# Test shuffling produces different order
loader_shuffle = DataLoader(dataset, batch_size=5, shuffle=True)
loader_no_shuffle = DataLoader(dataset, batch_size=5, shuffle=False)
batch_shuffle = list(loader_shuffle)[0]
batch_no_shuffle = list(loader_no_shuffle)[0]
# Note: This might occasionally fail due to random chance, but very unlikely
# We'll just test that both contain all the original data
shuffle_features = set(tuple(row) for row in batch_shuffle[0].data)
no_shuffle_features = set(tuple(row) for row in batch_no_shuffle[0].data)
expected_features = {(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)}
assert shuffle_features == expected_features, "Shuffle should preserve all data"
assert no_shuffle_features == expected_features, "No shuffle should preserve all data"
print("✅ DataLoader works correctly!")
if __name__ == "__main__":
test_unit_dataloader()In [ ]:
def analyze_dataloader_performance():
"""📊 Analyze DataLoader performance characteristics."""
print("📊 Analyzing DataLoader Performance...")
import time
# Create test dataset of varying sizes
sizes = [1000, 5000, 10000]
batch_sizes = [16, 64, 256]
print("\n🔍 Batch Size vs Loading Time:")
for size in sizes:
# Create synthetic dataset
features = Tensor(np.random.randn(size, 100)) # 100 features
labels = Tensor(np.random.randint(0, 10, size))
dataset = TensorDataset(features, labels)
print(f"\nDataset size: {size} samples")
for batch_size in batch_sizes:
# Time data loading
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
start_time = time.time()
batch_count = 0
for batch in loader:
batch_count += 1
end_time = time.time()
elapsed = end_time - start_time
throughput = size / elapsed if elapsed > 0 else float('inf')
print(f" Batch size {batch_size:3d}: {elapsed:.3f}s ({throughput:,.0f} samples/sec)")
# Analyze shuffle overhead
print("\n🔄 Shuffle Overhead Analysis:")
dataset_size = 10000
features = Tensor(np.random.randn(dataset_size, 50))
labels = Tensor(np.random.randint(0, 5, dataset_size))
dataset = TensorDataset(features, labels)
batch_size = 64
# No shuffle
loader_no_shuffle = DataLoader(dataset, batch_size=batch_size, shuffle=False)
start_time = time.time()
batches_no_shuffle = list(loader_no_shuffle)
time_no_shuffle = time.time() - start_time
# With shuffle
loader_shuffle = DataLoader(dataset, batch_size=batch_size, shuffle=True)
start_time = time.time()
batches_shuffle = list(loader_shuffle)
time_shuffle = time.time() - start_time
shuffle_overhead = ((time_shuffle - time_no_shuffle) / time_no_shuffle) * 100
print(f" No shuffle: {time_no_shuffle:.3f}s")
print(f" With shuffle: {time_shuffle:.3f}s")
print(f" Shuffle overhead: {shuffle_overhead:.1f}%")
print("\n💡 Key Insights:")
print("• Larger batch sizes reduce per-sample overhead")
print("• Shuffle adds minimal overhead for reasonable dataset sizes")
print("• Memory usage scales linearly with batch size")
print("🚀 Production tip: Balance batch size with GPU memory limits")
# analyze_dataloader_performance() # Optional: Run manually for performance insights
def analyze_memory_usage():
"""📊 Analyze memory usage patterns in data loading."""
print("\n📊 Analyzing Memory Usage Patterns...")
# Memory usage estimation
def estimate_memory_mb(batch_size, feature_size, dtype_bytes=4):
"""Estimate memory usage for a batch."""
return (batch_size * feature_size * dtype_bytes) / (1024 * 1024)
print("\n💾 Memory Usage by Batch Configuration:")
feature_sizes = [784, 3072, 50176] # MNIST, CIFAR-10, ImageNet-like
feature_names = ["MNIST (28×28)", "CIFAR-10 (32×32×3)", "ImageNet (224×224×1)"]
batch_sizes = [1, 32, 128, 512]
for feature_size, name in zip(feature_sizes, feature_names):
print(f"\n{name}:")
for batch_size in batch_sizes:
memory_mb = estimate_memory_mb(batch_size, feature_size)
print(f" Batch {batch_size:3d}: {memory_mb:6.1f} MB")
print("\n🎯 Memory Trade-offs:")
print("• Larger batches: More memory, better GPU utilization")
print("• Smaller batches: Less memory, more noisy gradients")
print("• Sweet spot: Usually 32-128 depending on model size")
# Demonstrate actual memory usage with our tensors
print("\n🔬 Actual Tensor Memory Usage:")
# Create different sized tensors
tensor_small = Tensor(np.random.randn(32, 784)) # Small batch
tensor_large = Tensor(np.random.randn(512, 784)) # Large batch
# Size in bytes (roughly)
small_bytes = tensor_small.data.nbytes
large_bytes = tensor_large.data.nbytes
print(f" Small batch (32×784): {small_bytes / 1024:.1f} KB")
print(f" Large batch (512×784): {large_bytes / 1024:.1f} KB")
print(f" Ratio: {large_bytes / small_bytes:.1f}×")
# analyze_memory_usage() # Optional: Run manually for memory insightsIn [ ]:
def test_training_integration():
"""🔬 Test DataLoader integration with training workflow."""
print("🔬 Integration Test: Training Workflow...")
# Create a realistic dataset
num_samples = 1000
num_features = 20
num_classes = 5
# Synthetic classification data
features = Tensor(np.random.randn(num_samples, num_features))
labels = Tensor(np.random.randint(0, num_classes, num_samples))
dataset = TensorDataset(features, labels)
# Create train/val splits
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
# Manual split (in production, you'd use proper splitting utilities)
train_indices = list(range(train_size))
val_indices = list(range(train_size, len(dataset)))
# Create subset datasets
train_samples = [dataset[i] for i in train_indices]
val_samples = [dataset[i] for i in val_indices]
# Convert back to tensors for TensorDataset
train_features = Tensor(np.stack([sample[0].data for sample in train_samples]))
train_labels = Tensor(np.stack([sample[1].data for sample in train_samples]))
val_features = Tensor(np.stack([sample[0].data for sample in val_samples]))
val_labels = Tensor(np.stack([sample[1].data for sample in val_samples]))
train_dataset = TensorDataset(train_features, train_labels)
val_dataset = TensorDataset(val_features, val_labels)
# Create DataLoaders
batch_size = 32
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
print(f"📊 Dataset splits:")
print(f" Training: {len(train_dataset)} samples, {len(train_loader)} batches")
print(f" Validation: {len(val_dataset)} samples, {len(val_loader)} batches")
# Simulate training loop
print("\n🏃 Simulated Training Loop:")
epoch_samples = 0
batch_count = 0
for batch_idx, (batch_features, batch_labels) in enumerate(train_loader):
batch_count += 1
epoch_samples += len(batch_features.data)
# Simulate forward pass (just check shapes)
assert batch_features.data.shape[0] <= batch_size, "Batch size exceeded"
assert batch_features.data.shape[1] == num_features, "Wrong feature count"
assert len(batch_labels.data) == len(batch_features.data), "Mismatched batch sizes"
if batch_idx < 3: # Show first few batches
print(f" Batch {batch_idx + 1}: {batch_features.data.shape[0]} samples")
print(f" Total: {batch_count} batches, {epoch_samples} samples processed")
# Validate that all samples were seen
assert epoch_samples == len(train_dataset), f"Expected {len(train_dataset)}, processed {epoch_samples}"
print("✅ Training integration works correctly!")In [ ]:
def test_module():
"""
Comprehensive test of entire module functionality.
This final test runs before module summary to ensure:
- All unit tests pass
- Functions work together correctly
- Module is ready for integration with TinyTorch
"""
print("🧪 RUNNING MODULE INTEGRATION TEST")
print("=" * 50)
# Run all unit tests
print("Running unit tests...")
test_unit_dataset()
test_unit_tensordataset()
test_unit_dataloader()
print("\nRunning integration scenarios...")
# Test complete workflow
test_training_integration()
print("\n" + "=" * 50)
print("🎉 ALL TESTS PASSED! Module ready for export.")
print("Run: tito module complete 08")In [ ]:
# Run comprehensive module test
if __name__ == "__main__":
test_module()