ayjays132's picture
Upload 478 files
101858b verified
Raw History Blame Contribute Delete
2.34 kB
"""
Memory Optimization Module Test
Quick test to verify all components work correctly.
"""
import torch
import sys
import os
# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from memory_optimization import (
UnifiedMemoryManager,
MemoryOptimizationConfig,
get_shared_model,
get_tensor,
return_tensor,
get_memory_stats,
clear_memory
)
def test_memory_optimization():
"""Test memory optimization module"""
print("=" * 70)
print("Testing Memory Optimization Module")
print("=" * 70)
# Initialize config
config = MemoryOptimizationConfig(
use_shared_model=True,
device="cuda" if torch.cuda.is_available() else "cpu"
)
print(f"\n[CONFIG] Device: {config.device}")
print(f"[CONFIG] Shared Model: {config.use_shared_model}")
# Initialize manager
manager = UnifiedMemoryManager(config)
print("\n[OK] UnifiedMemoryManager initialized")
# Test tensor pooling
print("\n[TEST] Tensor Pooling...")
tensor1 = manager.get_tensor((10, 1024), dtype=torch.float32)
print(f" [OK] Created tensor: {tensor1.shape}, device: {tensor1.device}")
manager.return_tensor(tensor1)
print(" [OK] Returned tensor to pool")
tensor2 = manager.get_tensor((10, 1024), dtype=torch.float32)
print(f" [OK] Retrieved tensor from pool: {tensor2.shape}")
# Test shared model (if available)
print("\n[TEST] Shared Model...")
try:
# This will use shared Qwen model if available
model = manager.get_shared_model("Qwen/Qwen3-0.6B", "transformer")
print(f" [OK] Got shared model: {type(model).__name__}")
except Exception as e:
print(f" [WARN] Could not get shared model: {e}")
# Test memory stats
print("\n[TEST] Memory Stats...")
stats = manager.get_memory_stats()
print(f" [OK] Got memory stats: {len(stats)} categories")
# Test cleanup
print("\n[TEST] Memory Cleanup...")
manager.clear_all_memory()
print(" [OK] Memory cleared")
print("\n" + "=" * 70)
print("[SUCCESS] All tests passed!")
print("=" * 70)
if __name__ == "__main__":
test_memory_optimization()