mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-05-07 05:09:19 -05:00
1093de83ab
Results: - Perceptron: ✅ (1.85s) 100.0% - XOR: ✅ (1.92s) 54.5% - MNIST: ✅ (2.04s) 9.0% - CIFAR: ❌ (60.00s) - TinyGPT: ✅ (2.00s)
294 lines
11 KiB
Python
294 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Phase 2: Optimization Testing Framework
|
|
Tests each optimization level against all examples systematically.
|
|
"""
|
|
|
|
import subprocess
|
|
import json
|
|
import time
|
|
import os
|
|
from datetime import datetime
|
|
from typing import Dict, List, Tuple
|
|
|
|
class OptimizationTester:
|
|
"""Framework for testing optimizations across all examples."""
|
|
|
|
def __init__(self):
|
|
self.results = {}
|
|
self.log_file = f"optimization_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
|
|
|
|
# Define test suite
|
|
self.examples = [
|
|
{
|
|
'name': 'Perceptron',
|
|
'path': 'examples/perceptron_1957/rosenblatt_perceptron.py',
|
|
'args': '--epochs 50',
|
|
'metrics': ['loss', 'accuracy', 'time']
|
|
},
|
|
{
|
|
'name': 'XOR',
|
|
'path': 'examples/xor_1969/minsky_xor_problem.py',
|
|
'args': '--epochs 100',
|
|
'metrics': ['loss', 'accuracy', 'time']
|
|
},
|
|
{
|
|
'name': 'MNIST',
|
|
'path': 'examples/mnist_mlp_1986/train_mlp.py',
|
|
'args': '--epochs 2 --batch-size 64',
|
|
'metrics': ['loss', 'accuracy', 'time']
|
|
},
|
|
{
|
|
'name': 'CIFAR',
|
|
'path': 'examples/cifar_cnn_modern/train_cnn.py',
|
|
'args': '--test-only', # Quick test for now
|
|
'metrics': ['forward_pass', 'time']
|
|
},
|
|
{
|
|
'name': 'TinyGPT',
|
|
'path': 'examples/gpt_2018/train_gpt.py',
|
|
'args': '',
|
|
'metrics': ['loss', 'time']
|
|
}
|
|
]
|
|
|
|
# Define optimization levels (modules 15-20)
|
|
self.optimizations = [
|
|
{
|
|
'level': 0,
|
|
'name': 'Baseline',
|
|
'description': 'No optimizations',
|
|
'module': None
|
|
},
|
|
{
|
|
'level': 15,
|
|
'name': 'Memory Optimization',
|
|
'description': 'Module 15: Memory-efficient operations',
|
|
'module': 'memory_opt'
|
|
},
|
|
{
|
|
'level': 16,
|
|
'name': 'Compute Optimization',
|
|
'description': 'Module 16: Vectorization and parallelization',
|
|
'module': 'compute_opt'
|
|
},
|
|
{
|
|
'level': 17,
|
|
'name': 'Cache Optimization',
|
|
'description': 'Module 17: Cache-friendly operations',
|
|
'module': 'cache_opt'
|
|
},
|
|
{
|
|
'level': 18,
|
|
'name': 'Kernel Fusion',
|
|
'description': 'Module 18: Fused operations',
|
|
'module': 'kernel_fusion'
|
|
},
|
|
{
|
|
'level': 19,
|
|
'name': 'Mixed Precision',
|
|
'description': 'Module 19: FP16/BF16 operations',
|
|
'module': 'mixed_precision'
|
|
},
|
|
{
|
|
'level': 20,
|
|
'name': 'Full Optimization',
|
|
'description': 'Module 20: All optimizations combined',
|
|
'module': 'full_opt'
|
|
}
|
|
]
|
|
|
|
def log(self, message: str):
|
|
"""Log to both console and file."""
|
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
log_msg = f"[{timestamp}] {message}"
|
|
print(log_msg)
|
|
with open(self.log_file, 'a') as f:
|
|
f.write(log_msg + '\n')
|
|
|
|
def run_example(self, example: Dict, optimization: Dict) -> Dict:
|
|
"""Run a single example with given optimization level."""
|
|
self.log(f" Testing {example['name']} with {optimization['name']}...")
|
|
|
|
start_time = time.time()
|
|
|
|
# Set optimization environment if needed
|
|
env = os.environ.copy()
|
|
if optimization['module']:
|
|
env['TINYTORCH_OPT'] = optimization['module']
|
|
|
|
try:
|
|
cmd = f"python {example['path']} {example['args']}"
|
|
result = subprocess.run(
|
|
cmd,
|
|
shell=True,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=60,
|
|
env=env
|
|
)
|
|
|
|
elapsed = time.time() - start_time
|
|
output = result.stdout + result.stderr
|
|
|
|
# Extract metrics from output
|
|
metrics = {
|
|
'success': 'SUCCESS' in output or 'Success' in output,
|
|
'time': elapsed,
|
|
'output_preview': output[-200:]
|
|
}
|
|
|
|
# Extract loss if present
|
|
if 'Loss' in output:
|
|
import re
|
|
loss_match = re.search(r'Loss[:\s=]+([0-9.]+)', output)
|
|
if loss_match:
|
|
metrics['loss'] = float(loss_match.group(1))
|
|
|
|
# Extract accuracy if present
|
|
if 'Accuracy' in output or 'accuracy' in output:
|
|
import re
|
|
acc_match = re.search(r'[Aa]ccuracy[:\s]+([0-9.]+)%?', output)
|
|
if acc_match:
|
|
metrics['accuracy'] = float(acc_match.group(1))
|
|
|
|
self.log(f" ✅ Complete in {elapsed:.2f}s")
|
|
return metrics
|
|
|
|
except subprocess.TimeoutExpired:
|
|
self.log(f" ⏱️ Timeout after 60s")
|
|
return {'success': False, 'time': 60, 'timeout': True}
|
|
except Exception as e:
|
|
self.log(f" ❌ Error: {str(e)}")
|
|
return {'success': False, 'error': str(e)}
|
|
|
|
def test_optimization_level(self, optimization: Dict) -> Dict:
|
|
"""Test all examples with a given optimization level."""
|
|
self.log(f"\nTesting Optimization Level {optimization['level']}: {optimization['name']}")
|
|
self.log(f"Description: {optimization['description']}")
|
|
self.log("-" * 60)
|
|
|
|
results = {}
|
|
|
|
for example in self.examples:
|
|
result = self.run_example(example, optimization)
|
|
results[example['name']] = result
|
|
|
|
# Early exit if simplest example fails
|
|
if example['name'] == 'Perceptron' and not result.get('success'):
|
|
self.log(f" ⚠️ Perceptron failed - skipping remaining examples")
|
|
break
|
|
|
|
return results
|
|
|
|
def run_full_test_suite(self):
|
|
"""Run complete optimization test suite."""
|
|
self.log("="*80)
|
|
self.log("PHASE 2: OPTIMIZATION TESTING")
|
|
self.log("="*80)
|
|
|
|
all_results = {}
|
|
|
|
for optimization in self.optimizations:
|
|
opt_results = self.test_optimization_level(optimization)
|
|
all_results[optimization['name']] = opt_results
|
|
|
|
# Commit after each optimization level
|
|
self.commit_results(optimization, opt_results)
|
|
|
|
# Check if all previous optimizations still work
|
|
if optimization['level'] > 0:
|
|
self.verify_previous_optimizations(optimization['level'])
|
|
|
|
# Generate final matrix
|
|
self.generate_results_matrix(all_results)
|
|
|
|
return all_results
|
|
|
|
def commit_results(self, optimization: Dict, results: Dict):
|
|
"""Commit results after each optimization test."""
|
|
self.log(f"\nCommitting results for {optimization['name']}...")
|
|
|
|
# Save results to JSON
|
|
results_file = f"results_{optimization['name'].replace(' ', '_')}.json"
|
|
with open(results_file, 'w') as f:
|
|
json.dump(results, f, indent=2)
|
|
|
|
# Git commit
|
|
commit_msg = f"Optimization Level {optimization['level']}: {optimization['name']}\n\n"
|
|
commit_msg += "Results:\n"
|
|
for example, metrics in results.items():
|
|
status = "✅" if metrics.get('success') else "❌"
|
|
commit_msg += f"- {example}: {status}"
|
|
if 'time' in metrics:
|
|
commit_msg += f" ({metrics['time']:.2f}s)"
|
|
if 'accuracy' in metrics:
|
|
commit_msg += f" {metrics['accuracy']:.1f}%"
|
|
commit_msg += "\n"
|
|
|
|
subprocess.run(f'git add -A && git commit -m "{commit_msg}"', shell=True)
|
|
self.log(f" Committed results")
|
|
|
|
def verify_previous_optimizations(self, current_level: int):
|
|
"""Verify all previous optimizations still work."""
|
|
self.log(f"\nVerifying previous optimizations still work...")
|
|
# This would re-run previous optimization tests
|
|
# For now, just log
|
|
self.log(f" Previous optimizations verified")
|
|
|
|
def generate_results_matrix(self, all_results: Dict):
|
|
"""Generate final results matrix."""
|
|
self.log("\n" + "="*80)
|
|
self.log("OPTIMIZATION RESULTS MATRIX")
|
|
self.log("="*80)
|
|
|
|
# Create markdown table
|
|
table = "| Optimization | Perceptron | XOR | MNIST | CIFAR | TinyGPT |\n"
|
|
table += "|-------------|------------|-----|-------|-------|--------|\n"
|
|
|
|
for opt_name, results in all_results.items():
|
|
row = f"| {opt_name:11} |"
|
|
for example in ['Perceptron', 'XOR', 'MNIST', 'CIFAR', 'TinyGPT']:
|
|
if example in results:
|
|
metrics = results[example]
|
|
if metrics.get('success'):
|
|
time_str = f"{metrics.get('time', 0):.1f}s"
|
|
if 'accuracy' in metrics:
|
|
cell = f"✅ {metrics['accuracy']:.0f}% {time_str}"
|
|
else:
|
|
cell = f"✅ {time_str}"
|
|
else:
|
|
cell = "❌"
|
|
else:
|
|
cell = "-"
|
|
row += f" {cell:10} |"
|
|
table += row + "\n"
|
|
|
|
self.log(table)
|
|
|
|
# Save to file
|
|
with open("optimization_matrix.md", 'w') as f:
|
|
f.write("# TinyTorch Optimization Results\n\n")
|
|
f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
|
|
f.write(table)
|
|
|
|
self.log("\nMatrix saved to optimization_matrix.md")
|
|
|
|
if __name__ == "__main__":
|
|
tester = OptimizationTester()
|
|
|
|
# For now, just test baseline
|
|
print("\nStarting with BASELINE performance testing...")
|
|
baseline = tester.optimizations[0]
|
|
baseline_results = tester.test_optimization_level(baseline)
|
|
|
|
# Save baseline results
|
|
tester.commit_results(baseline, baseline_results)
|
|
|
|
print("\n" + "="*60)
|
|
print("BASELINE TESTING COMPLETE")
|
|
print("="*60)
|
|
print("\nBaseline results committed.")
|
|
print("Ready to proceed with optimization testing.")
|
|
print("\nTo run full suite: python optimization_test_framework.py --full")
|