mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-07-19 12:16:13 -05:00
Vijay Janapa Reddi
8be87d0add
PROBLEM: - nbdev requires #| export directive on EACH cell to export when using # %% markers - Cell markers inside class definitions split classes across multiple cells - Only partial classes were being exported to tinytorch package - Missing matmul, arithmetic operations, and activation classes in exports SOLUTION: 1. Removed # %% cell markers INSIDE class definitions (kept classes as single units) 2. Added #| export to imports cell at top of each module 3. Added #| export before each exportable class definition in all 20 modules 4. Added __call__ method to Sigmoid for functional usage 5. Fixed numpy import (moved to module level from __init__) MODULES FIXED: - 01_tensor: Tensor class with all operations (matmul, arithmetic, shape ops) - 02_activations: Sigmoid, ReLU, Tanh, GELU, Softmax classes - 03_layers: Linear, Dropout classes - 04_losses: MSELoss, CrossEntropyLoss, BinaryCrossEntropyLoss classes - 05_autograd: Function, AddBackward, MulBackward, MatmulBackward, SumBackward - 06_optimizers: Optimizer, SGD, Adam, AdamW classes - 07_training: CosineSchedule, Trainer classes - 08_dataloader: Dataset, TensorDataset, DataLoader classes - 09_spatial: Conv2d, MaxPool2d, AvgPool2d, SimpleCNN classes - 10-20: All exportable classes in remaining modules TESTING: - Test functions use 'if __name__ == "__main__"' guards - Tests run in notebooks but NOT on import - Rosenblatt Perceptron milestone working perfectly RESULT: ✅ All 20 modules export correctly ✅ Perceptron (1957) milestone functional ✅ Clean separation: development (modules/source) vs package (tinytorch)
76 KiB
76 KiB
In [ ]:
#| default_exp optimization.compression
#| export
import numpy as np
import copy
from typing import List, Dict, Any, Tuple, Optional
import time
# Import from previous modules
# Note: In the full package, these would be imports like:
# from tinytorch.core.tensor import Tensor
# from tinytorch.core.layers import Linear
# For development, we'll create minimal implementations
class Tensor:
"""Minimal Tensor class for compression development - imports from Module 01 in practice."""
def __init__(self, data, requires_grad=False):
self.data = np.array(data)
self.shape = self.data.shape
self.size = self.data.size
self.requires_grad = requires_grad
self.grad = None
def __add__(self, other):
if isinstance(other, Tensor):
return Tensor(self.data + other.data)
return Tensor(self.data + other)
def __mul__(self, other):
if isinstance(other, Tensor):
return Tensor(self.data * other.data)
return Tensor(self.data * other)
def matmul(self, other):
return Tensor(np.dot(self.data, other.data))
def abs(self):
return Tensor(np.abs(self.data))
def sum(self, axis=None):
return Tensor(self.data.sum(axis=axis))
def __repr__(self):
return f"Tensor(shape={self.shape})"
class Linear:
"""Minimal Linear layer for compression development - imports from Module 03 in practice."""
def __init__(self, in_features, out_features, bias=True):
self.in_features = in_features
self.out_features = out_features
# Initialize with He initialization
self.weight = Tensor(np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features))
self.bias = Tensor(np.zeros(out_features)) if bias else None
def forward(self, x):
output = x.matmul(self.weight)
if self.bias is not None:
output = output + self.bias
return output
def parameters(self):
params = [self.weight]
if self.bias is not None:
params.append(self.bias)
return params
class Sequential:
"""Minimal Sequential container for model compression."""
def __init__(self, *layers):
self.layers = list(layers)
def forward(self, x):
for layer in self.layers:
x = layer.forward(x)
return x
def parameters(self):
params = []
for layer in self.layers:
if hasattr(layer, 'parameters'):
params.extend(layer.parameters())
return paramsIn [ ]:
def measure_sparsity(model) -> float:
"""
Calculate the percentage of zero weights in a model.
TODO: Count zero weights and total weights across all layers
APPROACH:
1. Iterate through all model parameters
2. Count zeros using np.sum(weights == 0)
3. Count total parameters
4. Return percentage: zeros / total * 100
EXAMPLE:
>>> model = Sequential(Linear(10, 5), Linear(5, 2))
>>> sparsity = measure_sparsity(model)
>>> print(f"Model sparsity: {sparsity:.1f}%")
Model sparsity: 0.0% # Before pruning
HINT: Use np.sum() to count zeros efficiently
"""
### BEGIN SOLUTION
total_params = 0
zero_params = 0
for param in model.parameters():
total_params += param.size
zero_params += np.sum(param.data == 0)
if total_params == 0:
return 0.0
return (zero_params / total_params) * 100.0
### END SOLUTION
def test_unit_measure_sparsity():
"""🔬 Test sparsity measurement functionality."""
print("🔬 Unit Test: Measure Sparsity...")
# Test with dense model
model = Sequential(Linear(4, 3), Linear(3, 2))
initial_sparsity = measure_sparsity(model)
assert initial_sparsity == 0.0, f"Expected 0% sparsity, got {initial_sparsity}%"
# Test with manually sparse model
model.layers[0].weight.data[0, 0] = 0
model.layers[0].weight.data[1, 1] = 0
sparse_sparsity = measure_sparsity(model)
assert sparse_sparsity > 0, f"Expected >0% sparsity, got {sparse_sparsity}%"
print("✅ measure_sparsity works correctly!")
test_unit_measure_sparsity()In [ ]:
def magnitude_prune(model, sparsity=0.9):
"""
Remove weights with smallest magnitudes to achieve target sparsity.
TODO: Implement global magnitude-based pruning
APPROACH:
1. Collect all weights from the model
2. Calculate absolute values to get magnitudes
3. Find threshold at desired sparsity percentile
4. Set weights below threshold to zero (in-place)
EXAMPLE:
>>> model = Sequential(Linear(100, 50), Linear(50, 10))
>>> original_params = sum(p.size for p in model.parameters())
>>> magnitude_prune(model, sparsity=0.8)
>>> final_sparsity = measure_sparsity(model)
>>> print(f"Achieved {final_sparsity:.1f}% sparsity")
Achieved 80.0% sparsity
HINTS:
- Use np.percentile() to find threshold
- Modify model parameters in-place
- Consider only weight matrices, not biases
"""
### BEGIN SOLUTION
# Collect all weights (excluding biases)
all_weights = []
weight_params = []
for param in model.parameters():
# Skip biases (typically 1D)
if len(param.shape) > 1:
all_weights.extend(param.data.flatten())
weight_params.append(param)
if not all_weights:
return
# Calculate magnitude threshold
magnitudes = np.abs(all_weights)
threshold = np.percentile(magnitudes, sparsity * 100)
# Apply pruning to each weight parameter
for param in weight_params:
mask = np.abs(param.data) >= threshold
param.data = param.data * mask
### END SOLUTION
def test_unit_magnitude_prune():
"""🔬 Test magnitude-based pruning functionality."""
print("🔬 Unit Test: Magnitude Prune...")
# Create test model with known weights
model = Sequential(Linear(4, 3), Linear(3, 2))
# Set specific weight values for predictable testing
model.layers[0].weight.data = np.array([
[1.0, 2.0, 3.0],
[0.1, 0.2, 0.3],
[4.0, 5.0, 6.0],
[0.01, 0.02, 0.03]
])
initial_sparsity = measure_sparsity(model)
assert initial_sparsity == 0.0, "Model should start with no sparsity"
# Apply 50% pruning
magnitude_prune(model, sparsity=0.5)
final_sparsity = measure_sparsity(model)
# Should achieve approximately 50% sparsity
assert 40 <= final_sparsity <= 60, f"Expected ~50% sparsity, got {final_sparsity}%"
# Verify largest weights survived
remaining_weights = model.layers[0].weight.data[model.layers[0].weight.data != 0]
assert len(remaining_weights) > 0, "Some weights should remain"
assert np.all(np.abs(remaining_weights) >= 0.1), "Large weights should survive"
print("✅ magnitude_prune works correctly!")
test_unit_magnitude_prune()In [ ]:
def structured_prune(model, prune_ratio=0.5):
"""
Remove entire channels/neurons based on L2 norm importance.
TODO: Implement structured pruning for Linear layers
APPROACH:
1. For each Linear layer, calculate L2 norm of each output channel
2. Rank channels by importance (L2 norm)
3. Remove lowest importance channels by setting to zero
4. This creates block sparsity that's hardware-friendly
EXAMPLE:
>>> model = Sequential(Linear(100, 50), Linear(50, 10))
>>> original_shape = model.layers[0].weight.shape
>>> structured_prune(model, prune_ratio=0.3)
>>> # 30% of channels are now completely zero
>>> final_sparsity = measure_sparsity(model)
>>> print(f"Structured sparsity: {final_sparsity:.1f}%")
Structured sparsity: 30.0%
HINTS:
- Calculate L2 norm along input dimension for each output channel
- Use np.linalg.norm(weights[:, channel]) for channel importance
- Set entire channels to zero (not just individual weights)
"""
### BEGIN SOLUTION
for layer in model.layers:
if isinstance(layer, Linear) and hasattr(layer, 'weight'):
weight = layer.weight.data
# Calculate L2 norm for each output channel (column)
channel_norms = np.linalg.norm(weight, axis=0)
# Find channels to prune (lowest importance)
num_channels = weight.shape[1]
num_to_prune = int(num_channels * prune_ratio)
if num_to_prune > 0:
# Get indices of channels to prune (smallest norms)
prune_indices = np.argpartition(channel_norms, num_to_prune)[:num_to_prune]
# Zero out entire channels
weight[:, prune_indices] = 0
# Also zero corresponding bias elements if bias exists
if layer.bias is not None:
layer.bias.data[prune_indices] = 0
### END SOLUTION
def test_unit_structured_prune():
"""🔬 Test structured pruning functionality."""
print("🔬 Unit Test: Structured Prune...")
# Create test model
model = Sequential(Linear(4, 6), Linear(6, 2))
# Set predictable weights for testing
model.layers[0].weight.data = np.array([
[1.0, 0.1, 2.0, 0.05, 3.0, 0.01], # Channels with varying importance
[1.1, 0.11, 2.1, 0.06, 3.1, 0.02],
[1.2, 0.12, 2.2, 0.07, 3.2, 0.03],
[1.3, 0.13, 2.3, 0.08, 3.3, 0.04]
])
initial_sparsity = measure_sparsity(model)
assert initial_sparsity == 0.0, "Model should start with no sparsity"
# Apply 33% structured pruning (2 out of 6 channels)
structured_prune(model, prune_ratio=0.33)
final_sparsity = measure_sparsity(model)
# Check that some channels are completely zero
weight = model.layers[0].weight.data
zero_channels = np.sum(np.all(weight == 0, axis=0))
assert zero_channels >= 1, f"Expected at least 1 zero channel, got {zero_channels}"
# Check that non-zero channels are completely preserved
for col in range(weight.shape[1]):
channel = weight[:, col]
assert np.all(channel == 0) or np.all(channel != 0), "Channels should be fully zero or fully non-zero"
print("✅ structured_prune works correctly!")
test_unit_structured_prune()In [ ]:
def low_rank_approximate(weight_matrix, rank_ratio=0.5):
"""
Approximate weight matrix using low-rank decomposition (SVD).
TODO: Implement SVD-based low-rank approximation
APPROACH:
1. Perform SVD: W = U @ S @ V^T
2. Keep only top k singular values where k = rank_ratio * min(dimensions)
3. Reconstruct: W_approx = U[:,:k] @ diag(S[:k]) @ V[:k,:]
4. Return decomposed matrices for memory savings
EXAMPLE:
>>> weight = np.random.randn(100, 50)
>>> U, S, V = low_rank_approximate(weight, rank_ratio=0.3)
>>> # Original: 100*50 = 5000 params
>>> # Compressed: 100*15 + 15*50 = 2250 params (55% reduction)
HINTS:
- Use np.linalg.svd() for decomposition
- Choose k = int(rank_ratio * min(m, n))
- Return U[:,:k], S[:k], V[:k,:] for reconstruction
"""
### BEGIN SOLUTION
m, n = weight_matrix.shape
# Perform SVD
U, S, V = np.linalg.svd(weight_matrix, full_matrices=False)
# Determine target rank
max_rank = min(m, n)
target_rank = max(1, int(rank_ratio * max_rank))
# Truncate to target rank
U_truncated = U[:, :target_rank]
S_truncated = S[:target_rank]
V_truncated = V[:target_rank, :]
return U_truncated, S_truncated, V_truncated
### END SOLUTION
def test_unit_low_rank_approximate():
"""🔬 Test low-rank approximation functionality."""
print("🔬 Unit Test: Low-Rank Approximate...")
# Create test weight matrix
original_weight = np.random.randn(20, 15)
original_params = original_weight.size
# Apply low-rank approximation
U, S, V = low_rank_approximate(original_weight, rank_ratio=0.4)
# Check dimensions
target_rank = int(0.4 * min(20, 15)) # min(20,15) = 15, so 0.4*15 = 6
assert U.shape == (20, target_rank), f"Expected U shape (20, {target_rank}), got {U.shape}"
assert S.shape == (target_rank,), f"Expected S shape ({target_rank},), got {S.shape}"
assert V.shape == (target_rank, 15), f"Expected V shape ({target_rank}, 15), got {V.shape}"
# Check parameter reduction
compressed_params = U.size + S.size + V.size
compression_ratio = compressed_params / original_params
assert compression_ratio < 1.0, f"Should compress, but ratio is {compression_ratio}"
# Check reconstruction quality
reconstructed = U @ np.diag(S) @ V
reconstruction_error = np.linalg.norm(original_weight - reconstructed)
relative_error = reconstruction_error / np.linalg.norm(original_weight)
assert relative_error < 0.5, f"Reconstruction error too high: {relative_error}"
print("✅ low_rank_approximate works correctly!")
test_unit_low_rank_approximate()In [ ]:
class KnowledgeDistillation:
"""
Knowledge distillation for model compression.
Train a smaller student model to mimic a larger teacher model.
"""
def __init__(self, teacher_model, student_model, temperature=3.0, alpha=0.7):
"""
Initialize knowledge distillation.
TODO: Set up teacher and student models with distillation parameters
APPROACH:
1. Store teacher and student models
2. Set temperature for softening probability distributions
3. Set alpha for balancing hard vs soft targets
Args:
teacher_model: Large, pre-trained model
student_model: Smaller model to train
temperature: Softening parameter for distributions
alpha: Weight for soft target loss (1-alpha for hard targets)
"""
### BEGIN SOLUTION
self.teacher_model = teacher_model
self.student_model = student_model
self.temperature = temperature
self.alpha = alpha
### END SOLUTION
def distillation_loss(self, student_logits, teacher_logits, true_labels):
"""
Calculate combined distillation loss.
TODO: Implement knowledge distillation loss function
APPROACH:
1. Calculate hard target loss (student vs true labels)
2. Calculate soft target loss (student vs teacher, with temperature)
3. Combine losses: alpha * soft_loss + (1-alpha) * hard_loss
EXAMPLE:
>>> kd = KnowledgeDistillation(teacher, student)
>>> loss = kd.distillation_loss(student_out, teacher_out, labels)
>>> print(f"Distillation loss: {loss:.4f}")
HINTS:
- Use temperature to soften distributions: logits/temperature
- Soft targets use KL divergence or cross-entropy
- Hard targets use standard classification loss
"""
### BEGIN SOLUTION
# Convert to numpy for this implementation
if hasattr(student_logits, 'data'):
student_logits = student_logits.data
if hasattr(teacher_logits, 'data'):
teacher_logits = teacher_logits.data
if hasattr(true_labels, 'data'):
true_labels = true_labels.data
# Soften distributions with temperature
student_soft = self._softmax(student_logits / self.temperature)
teacher_soft = self._softmax(teacher_logits / self.temperature)
# Soft target loss (KL divergence)
soft_loss = self._kl_divergence(student_soft, teacher_soft)
# Hard target loss (cross-entropy)
student_hard = self._softmax(student_logits)
hard_loss = self._cross_entropy(student_hard, true_labels)
# Combined loss
total_loss = self.alpha * soft_loss + (1 - self.alpha) * hard_loss
return total_loss
### END SOLUTION
def _softmax(self, logits):
"""Compute softmax with numerical stability."""
exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
return exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
def _kl_divergence(self, p, q):
"""Compute KL divergence between distributions."""
return np.sum(p * np.log(p / (q + 1e-8) + 1e-8))
def _cross_entropy(self, predictions, labels):
"""Compute cross-entropy loss."""
# Simple implementation for integer labels
if labels.ndim == 1:
return -np.mean(np.log(predictions[np.arange(len(labels)), labels] + 1e-8))
else:
return -np.mean(np.sum(labels * np.log(predictions + 1e-8), axis=1))
def test_unit_knowledge_distillation():
"""🔬 Test knowledge distillation functionality."""
print("🔬 Unit Test: Knowledge Distillation...")
# Create teacher and student models
teacher = Sequential(Linear(10, 20), Linear(20, 5))
student = Sequential(Linear(10, 5)) # Smaller model
# Initialize knowledge distillation
kd = KnowledgeDistillation(teacher, student, temperature=3.0, alpha=0.7)
# Create dummy data
input_data = Tensor(np.random.randn(8, 10)) # Batch of 8
true_labels = np.array([0, 1, 2, 3, 4, 0, 1, 2]) # Class labels
# Forward passes
teacher_output = teacher.forward(input_data)
student_output = student.forward(input_data)
# Calculate distillation loss
loss = kd.distillation_loss(student_output, teacher_output, true_labels)
# Verify loss is reasonable
assert isinstance(loss, (float, np.floating)), f"Loss should be float, got {type(loss)}"
assert loss > 0, f"Loss should be positive, got {loss}"
assert not np.isnan(loss), "Loss should not be NaN"
print("✅ knowledge_distillation works correctly!")
test_unit_knowledge_distillation()In [ ]:
def compress_model(model, compression_config):
"""
Apply comprehensive model compression based on configuration.
TODO: Implement complete compression pipeline
APPROACH:
1. Apply magnitude pruning if specified
2. Apply structured pruning if specified
3. Apply low-rank approximation if specified
4. Return compression statistics
EXAMPLE:
>>> config = {
... 'magnitude_prune': 0.8,
... 'structured_prune': 0.3,
... 'low_rank': 0.5
... }
>>> stats = compress_model(model, config)
>>> print(f"Final sparsity: {stats['sparsity']:.1f}%")
Final sparsity: 85.0%
HINT: Apply techniques sequentially and measure results
"""
### BEGIN SOLUTION
original_params = sum(p.size for p in model.parameters())
original_sparsity = measure_sparsity(model)
stats = {
'original_params': original_params,
'original_sparsity': original_sparsity,
'applied_techniques': []
}
# Apply magnitude pruning
if 'magnitude_prune' in compression_config:
sparsity = compression_config['magnitude_prune']
magnitude_prune(model, sparsity=sparsity)
stats['applied_techniques'].append(f'magnitude_prune_{sparsity}')
# Apply structured pruning
if 'structured_prune' in compression_config:
ratio = compression_config['structured_prune']
structured_prune(model, prune_ratio=ratio)
stats['applied_techniques'].append(f'structured_prune_{ratio}')
# Apply low-rank approximation (conceptually - would need architecture changes)
if 'low_rank' in compression_config:
ratio = compression_config['low_rank']
# For demo, we'll just record that it would be applied
stats['applied_techniques'].append(f'low_rank_{ratio}')
# Final measurements
final_sparsity = measure_sparsity(model)
stats['final_sparsity'] = final_sparsity
stats['sparsity_increase'] = final_sparsity - original_sparsity
return stats
### END SOLUTION
def test_unit_compress_model():
"""🔬 Test comprehensive model compression."""
print("🔬 Unit Test: Compress Model...")
# Create test model
model = Sequential(Linear(20, 15), Linear(15, 10), Linear(10, 5))
# Define compression configuration
config = {
'magnitude_prune': 0.7,
'structured_prune': 0.2
}
# Apply compression
stats = compress_model(model, config)
# Verify statistics
assert 'original_params' in stats, "Should track original parameter count"
assert 'final_sparsity' in stats, "Should track final sparsity"
assert 'applied_techniques' in stats, "Should track applied techniques"
# Verify compression was applied
assert stats['final_sparsity'] > stats['original_sparsity'], "Sparsity should increase"
assert len(stats['applied_techniques']) == 2, "Should apply both techniques"
# Verify model still has reasonable structure
remaining_params = sum(np.count_nonzero(p.data) for p in model.parameters())
assert remaining_params > 0, "Model should retain some parameters"
print("✅ compress_model works correctly!")
test_unit_compress_model()In [ ]:
def analyze_compression_ratios():
"""📊 Analyze compression ratios for different techniques."""
print("📊 Analyzing Compression Ratios...")
# Create test models of different sizes
models = {
'Small': Sequential(Linear(50, 30), Linear(30, 10)),
'Medium': Sequential(Linear(200, 128), Linear(128, 64), Linear(64, 10)),
'Large': Sequential(Linear(500, 256), Linear(256, 128), Linear(128, 10))
}
compression_techniques = [
('Magnitude 50%', {'magnitude_prune': 0.5}),
('Magnitude 90%', {'magnitude_prune': 0.9}),
('Structured 30%', {'structured_prune': 0.3}),
('Combined', {'magnitude_prune': 0.8, 'structured_prune': 0.2})
]
print(f"{'Model':<8} {'Technique':<15} {'Original':<10} {'Final':<10} {'Reduction':<10}")
print("-" * 65)
for model_name, model in models.items():
original_params = sum(p.size for p in model.parameters())
for tech_name, config in compression_techniques:
# Create fresh copy for each test
test_model = copy.deepcopy(model)
# Apply compression
stats = compress_model(test_model, config)
# Calculate compression ratio
remaining_params = sum(np.count_nonzero(p.data) for p in test_model.parameters())
reduction = (1 - remaining_params / original_params) * 100
print(f"{model_name:<8} {tech_name:<15} {original_params:<10} {remaining_params:<10} {reduction:<9.1f}%")
print("\n💡 Key Insights:")
print("• Magnitude pruning achieves predictable sparsity levels")
print("• Structured pruning creates hardware-friendly sparsity")
print("• Combined techniques offer maximum compression")
print("• Larger models compress better (more redundancy)")
analyze_compression_ratios()In [ ]:
def analyze_compression_speed():
"""📊 Analyze inference speed with different compression levels."""
print("📊 Analyzing Compression Speed Impact...")
# Create test model
model = Sequential(Linear(512, 256), Linear(256, 128), Linear(128, 10))
test_input = Tensor(np.random.randn(100, 512)) # Batch of 100
def time_inference(model, input_data, iterations=50):
"""Time model inference."""
times = []
for _ in range(iterations):
start = time.time()
_ = model.forward(input_data)
times.append(time.time() - start)
return np.mean(times[5:]) # Skip first few for warmup
# Test different compression levels
compression_levels = [
('Original', {}),
('Light Pruning', {'magnitude_prune': 0.5}),
('Heavy Pruning', {'magnitude_prune': 0.9}),
('Structured', {'structured_prune': 0.3}),
('Combined', {'magnitude_prune': 0.8, 'structured_prune': 0.2})
]
print(f"{'Compression':<15} {'Sparsity':<10} {'Time (ms)':<12} {'Speedup':<10}")
print("-" * 50)
baseline_time = None
for name, config in compression_levels:
# Create fresh model copy
test_model = copy.deepcopy(model)
# Apply compression
if config:
compress_model(test_model, config)
# Measure performance
sparsity = measure_sparsity(test_model)
inference_time = time_inference(test_model, test_input) * 1000 # Convert to ms
if baseline_time is None:
baseline_time = inference_time
speedup = 1.0
else:
speedup = baseline_time / inference_time
print(f"{name:<15} {sparsity:<9.1f}% {inference_time:<11.2f} {speedup:<9.2f}x")
print("\n💡 Speed Insights:")
print("• Dense matrix operations show minimal speedup from unstructured sparsity")
print("• Structured sparsity enables better hardware acceleration")
print("• Real speedups require sparse-optimized libraries (e.g., NVIDIA 2:4 sparsity)")
print("• Memory bandwidth often more important than parameter count")
analyze_compression_speed()In [ ]:
def analyze_compression_accuracy_tradeoff():
"""📊 Analyze accuracy vs compression trade-offs."""
print("📊 Analyzing Accuracy vs Compression Trade-offs...")
# Simulate accuracy degradation (in practice, would need real training/testing)
def simulate_accuracy_loss(sparsity, technique_type):
"""Simulate realistic accuracy loss patterns."""
if technique_type == 'magnitude':
# Magnitude pruning: gradual degradation
return max(0, sparsity * 0.3 + np.random.normal(0, 0.05))
elif technique_type == 'structured':
# Structured pruning: more aggressive early loss
return max(0, sparsity * 0.5 + np.random.normal(0, 0.1))
elif technique_type == 'knowledge_distillation':
# Knowledge distillation: better preservation
return max(0, sparsity * 0.1 + np.random.normal(0, 0.02))
else:
return sparsity * 0.4
# Test different compression strategies
strategies = [
('Magnitude Only', 'magnitude'),
('Structured Only', 'structured'),
('Knowledge Distillation', 'knowledge_distillation'),
('Combined Approach', 'combined')
]
sparsity_levels = np.arange(0.1, 1.0, 0.1)
print(f"{'Strategy':<20} {'Sparsity':<10} {'Accuracy Loss':<15}")
print("-" * 50)
for strategy_name, strategy_type in strategies:
print(f"\n{strategy_name}:")
for sparsity in sparsity_levels:
if strategy_type == 'combined':
# Combined approach uses multiple techniques
loss = min(
simulate_accuracy_loss(sparsity * 0.7, 'magnitude'),
simulate_accuracy_loss(sparsity * 0.3, 'structured')
)
else:
loss = simulate_accuracy_loss(sparsity, strategy_type)
print(f"{'':20} {sparsity:<9.1f} {loss:<14.3f}")
print("\n💡 Trade-off Insights:")
print("• Knowledge distillation preserves accuracy best at high compression")
print("• Magnitude pruning offers gradual degradation curve")
print("• Structured pruning enables hardware acceleration but higher accuracy loss")
print("• Combined approaches balance multiple objectives")
print("• Early stopping based on accuracy threshold is crucial")
analyze_compression_accuracy_tradeoff()In [ ]:
def test_module():
"""
Comprehensive test of entire compression 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_measure_sparsity()
test_unit_magnitude_prune()
test_unit_structured_prune()
test_unit_low_rank_approximate()
test_unit_knowledge_distillation()
test_unit_compress_model()
print("\nRunning integration scenarios...")
# Test 1: Complete compression pipeline
print("🔬 Integration Test: Complete compression pipeline...")
# Create a realistic model
model = Sequential(
Linear(784, 512), # Input layer (like MNIST)
Linear(512, 256), # Hidden layer 1
Linear(256, 128), # Hidden layer 2
Linear(128, 10) # Output layer
)
original_params = sum(p.size for p in model.parameters())
print(f"Original model: {original_params:,} parameters")
# Apply comprehensive compression
compression_config = {
'magnitude_prune': 0.8,
'structured_prune': 0.3
}
stats = compress_model(model, compression_config)
final_sparsity = measure_sparsity(model)
# Validate compression results
assert final_sparsity > 70, f"Expected >70% sparsity, got {final_sparsity:.1f}%"
assert stats['sparsity_increase'] > 70, "Should achieve significant compression"
assert len(stats['applied_techniques']) == 2, "Should apply both techniques"
print(f"✅ Achieved {final_sparsity:.1f}% sparsity with {len(stats['applied_techniques'])} techniques")
# Test 2: Knowledge distillation setup
print("🔬 Integration Test: Knowledge distillation...")
teacher = Sequential(Linear(100, 200), Linear(200, 50))
student = Sequential(Linear(100, 50)) # 3x fewer parameters
kd = KnowledgeDistillation(teacher, student, temperature=4.0, alpha=0.8)
# Verify setup
teacher_params = sum(p.size for p in teacher.parameters())
student_params = sum(p.size for p in student.parameters())
compression_ratio = student_params / teacher_params
assert compression_ratio < 0.5, f"Student should be <50% of teacher size, got {compression_ratio:.2f}"
assert kd.temperature == 4.0, "Temperature should be set correctly"
assert kd.alpha == 0.8, "Alpha should be set correctly"
print(f"✅ Knowledge distillation: {compression_ratio:.2f}x size reduction")
# Test 3: Low-rank approximation
print("🔬 Integration Test: Low-rank approximation...")
large_matrix = np.random.randn(200, 150)
U, S, V = low_rank_approximate(large_matrix, rank_ratio=0.3)
original_size = large_matrix.size
compressed_size = U.size + S.size + V.size
compression_ratio = compressed_size / original_size
assert compression_ratio < 0.7, f"Should achieve compression, got ratio {compression_ratio:.2f}"
# Test reconstruction
reconstructed = U @ np.diag(S) @ V
error = np.linalg.norm(large_matrix - reconstructed) / np.linalg.norm(large_matrix)
assert error < 0.5, f"Reconstruction error too high: {error:.3f}"
print(f"✅ Low-rank: {compression_ratio:.2f}x compression, {error:.3f} error")
print("\n" + "=" * 50)
print("🎉 ALL TESTS PASSED! Module ready for export.")
print("Run: tito module complete 18")
# Call the integration test
test_module()In [ ]:
if __name__ == "__main__":
print("🚀 Running Compression module...")
test_module()
print("✅ Module validation complete!")