371 lines
11 KiB
Python
371 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Comprehensive NPU Integration Test for Strix Halo
|
|
Tests NPU acceleration with your trading models
|
|
"""
|
|
import sys
|
|
import os
|
|
import time
|
|
import logging
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
# Add project root to path
|
|
sys.path.append('/mnt/shared/DEV/repos/d-popov.com/gogo2')
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def test_npu_detection():
|
|
"""Test NPU detection and setup"""
|
|
print("=== NPU Detection Test ===")
|
|
|
|
try:
|
|
from utils.npu_detector import get_npu_info, is_npu_available, get_onnx_providers
|
|
|
|
info = get_npu_info()
|
|
print(f"NPU Available: {info['available']}")
|
|
print(f"NPU Info: {info['info']}")
|
|
|
|
providers = get_onnx_providers()
|
|
print(f"ONNX Providers: {providers}")
|
|
|
|
if is_npu_available():
|
|
print("✅ NPU is available!")
|
|
return True
|
|
else:
|
|
print("❌ NPU not available")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"❌ NPU detection failed: {e}")
|
|
return False
|
|
|
|
def test_onnx_runtime():
|
|
"""Test ONNX Runtime functionality"""
|
|
print("\n=== ONNX Runtime Test ===")
|
|
|
|
try:
|
|
import onnxruntime as ort
|
|
print(f"ONNX Runtime version: {ort.__version__}")
|
|
|
|
# Test providers
|
|
providers = ort.get_available_providers()
|
|
print(f"Available providers: {providers}")
|
|
|
|
# Test DirectML provider
|
|
if 'DmlExecutionProvider' in providers:
|
|
print("✅ DirectML provider available")
|
|
else:
|
|
print("❌ DirectML provider not available")
|
|
|
|
return True
|
|
|
|
except ImportError:
|
|
print("❌ ONNX Runtime not installed")
|
|
return False
|
|
except Exception as e:
|
|
print(f"❌ ONNX Runtime test failed: {e}")
|
|
return False
|
|
|
|
def create_test_model():
|
|
"""Create a simple test model for NPU testing"""
|
|
class SimpleTradingModel(nn.Module):
|
|
def __init__(self, input_size=50, hidden_size=128, output_size=3):
|
|
super().__init__()
|
|
self.fc1 = nn.Linear(input_size, hidden_size)
|
|
self.fc2 = nn.Linear(hidden_size, hidden_size)
|
|
self.fc3 = nn.Linear(hidden_size, output_size)
|
|
self.relu = nn.ReLU()
|
|
self.dropout = nn.Dropout(0.1)
|
|
|
|
def forward(self, x):
|
|
x = self.relu(self.fc1(x))
|
|
x = self.dropout(x)
|
|
x = self.relu(self.fc2(x))
|
|
x = self.dropout(x)
|
|
x = self.fc3(x)
|
|
return x
|
|
|
|
return SimpleTradingModel()
|
|
|
|
def test_model_conversion():
|
|
"""Test PyTorch to ONNX conversion"""
|
|
print("\n=== Model Conversion Test ===")
|
|
|
|
try:
|
|
from utils.npu_acceleration import PyTorchToONNXConverter
|
|
|
|
# Create test model
|
|
model = create_test_model()
|
|
model.eval()
|
|
|
|
# Create converter
|
|
converter = PyTorchToONNXConverter(model)
|
|
|
|
# Convert to ONNX
|
|
onnx_path = "/tmp/test_trading_model.onnx"
|
|
input_shape = (50,) # 50 features
|
|
|
|
success = converter.convert(
|
|
output_path=onnx_path,
|
|
input_shape=input_shape,
|
|
input_names=['trading_features'],
|
|
output_names=['trading_signals']
|
|
)
|
|
|
|
if success:
|
|
print("✅ Model conversion successful")
|
|
|
|
# Verify the model
|
|
if converter.verify_onnx_model(onnx_path, input_shape):
|
|
print("✅ ONNX model verification successful")
|
|
return True
|
|
else:
|
|
print("❌ ONNX model verification failed")
|
|
return False
|
|
else:
|
|
print("❌ Model conversion failed")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"❌ Model conversion test failed: {e}")
|
|
return False
|
|
|
|
def test_npu_acceleration():
|
|
"""Test NPU-accelerated inference"""
|
|
print("\n=== NPU Acceleration Test ===")
|
|
|
|
try:
|
|
from utils.npu_acceleration import NPUAcceleratedModel
|
|
|
|
# Create test model
|
|
model = create_test_model()
|
|
model.eval()
|
|
|
|
# Create NPU-accelerated model
|
|
npu_model = NPUAcceleratedModel(
|
|
pytorch_model=model,
|
|
model_name="test_trading_model",
|
|
input_shape=(50,)
|
|
)
|
|
|
|
# Test inference
|
|
test_input = np.random.randn(1, 50).astype(np.float32)
|
|
|
|
start_time = time.time()
|
|
output = npu_model.predict(test_input)
|
|
inference_time = (time.time() - start_time) * 1000 # ms
|
|
|
|
print(f"✅ NPU inference successful")
|
|
print(f"Inference time: {inference_time:.2f} ms")
|
|
print(f"Output shape: {output.shape}")
|
|
|
|
# Get performance info
|
|
perf_info = npu_model.get_performance_info()
|
|
print(f"Performance info: {perf_info}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ NPU acceleration test failed: {e}")
|
|
return False
|
|
|
|
def test_model_interfaces():
|
|
"""Test enhanced model interfaces with NPU support"""
|
|
print("\n=== Model Interfaces Test ===")
|
|
|
|
try:
|
|
from NN.models.model_interfaces import CNNModelInterface, RLAgentInterface
|
|
|
|
# Create test models
|
|
cnn_model = create_test_model()
|
|
rl_model = create_test_model()
|
|
|
|
# Test CNN interface
|
|
cnn_interface = CNNModelInterface(
|
|
model=cnn_model,
|
|
name="test_cnn",
|
|
enable_npu=True,
|
|
input_shape=(50,)
|
|
)
|
|
|
|
# Test RL interface
|
|
rl_interface = RLAgentInterface(
|
|
model=rl_model,
|
|
name="test_rl",
|
|
enable_npu=True,
|
|
input_shape=(50,)
|
|
)
|
|
|
|
# Test predictions
|
|
test_data = np.random.randn(1, 50).astype(np.float32)
|
|
|
|
cnn_output = cnn_interface.predict(test_data)
|
|
rl_output = rl_interface.predict(test_data)
|
|
|
|
print(f"✅ CNN interface prediction: {cnn_output is not None}")
|
|
print(f"✅ RL interface prediction: {rl_output is not None}")
|
|
|
|
# Test acceleration info
|
|
cnn_info = cnn_interface.get_acceleration_info()
|
|
rl_info = rl_interface.get_acceleration_info()
|
|
|
|
print(f"CNN acceleration info: {cnn_info}")
|
|
print(f"RL acceleration info: {rl_info}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ Model interfaces test failed: {e}")
|
|
return False
|
|
|
|
def benchmark_performance():
|
|
"""Benchmark NPU vs CPU performance"""
|
|
print("\n=== Performance Benchmark ===")
|
|
|
|
try:
|
|
from utils.npu_acceleration import NPUAcceleratedModel
|
|
|
|
# Create test model
|
|
model = create_test_model()
|
|
model.eval()
|
|
|
|
# Create NPU-accelerated model
|
|
npu_model = NPUAcceleratedModel(
|
|
pytorch_model=model,
|
|
model_name="benchmark_model",
|
|
input_shape=(50,)
|
|
)
|
|
|
|
# Test data
|
|
test_data = np.random.randn(100, 50).astype(np.float32)
|
|
|
|
# Benchmark NPU inference
|
|
if npu_model.onnx_model:
|
|
npu_times = []
|
|
for i in range(10):
|
|
start_time = time.time()
|
|
npu_model.predict(test_data[i:i+1])
|
|
npu_times.append((time.time() - start_time) * 1000)
|
|
|
|
avg_npu_time = np.mean(npu_times)
|
|
print(f"Average NPU inference time: {avg_npu_time:.2f} ms")
|
|
|
|
# Benchmark CPU inference
|
|
cpu_times = []
|
|
model.eval()
|
|
with torch.no_grad():
|
|
for i in range(10):
|
|
start_time = time.time()
|
|
input_tensor = torch.from_numpy(test_data[i:i+1])
|
|
model(input_tensor)
|
|
cpu_times.append((time.time() - start_time) * 1000)
|
|
|
|
avg_cpu_time = np.mean(cpu_times)
|
|
print(f"Average CPU inference time: {avg_cpu_time:.2f} ms")
|
|
|
|
if npu_model.onnx_model:
|
|
speedup = avg_cpu_time / avg_npu_time
|
|
print(f"NPU speedup: {speedup:.2f}x")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ Performance benchmark failed: {e}")
|
|
return False
|
|
|
|
def test_integration_with_existing_models():
|
|
"""Test integration with existing trading models"""
|
|
print("\n=== Integration Test ===")
|
|
|
|
try:
|
|
# Test with existing CNN model
|
|
from NN.models.cnn_model import EnhancedCNNModel
|
|
|
|
# Create a small CNN model for testing
|
|
cnn_model = EnhancedCNNModel(
|
|
input_size=60,
|
|
feature_dim=50,
|
|
output_size=3
|
|
)
|
|
|
|
# Test NPU acceleration
|
|
from utils.npu_acceleration import NPUAcceleratedModel
|
|
|
|
npu_cnn = NPUAcceleratedModel(
|
|
pytorch_model=cnn_model,
|
|
model_name="enhanced_cnn_test",
|
|
input_shape=(60, 50)
|
|
)
|
|
|
|
# Test inference
|
|
test_input = np.random.randn(1, 60, 50).astype(np.float32)
|
|
output = npu_cnn.predict(test_input)
|
|
|
|
print(f"✅ Enhanced CNN NPU integration successful")
|
|
print(f"Output shape: {output.shape}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ Integration test failed: {e}")
|
|
return False
|
|
|
|
def main():
|
|
"""Run all NPU tests"""
|
|
print("Starting Strix Halo NPU Integration Tests...")
|
|
print("=" * 50)
|
|
|
|
tests = [
|
|
("NPU Detection", test_npu_detection),
|
|
("ONNX Runtime", test_onnx_runtime),
|
|
("Model Conversion", test_model_conversion),
|
|
("NPU Acceleration", test_npu_acceleration),
|
|
("Model Interfaces", test_model_interfaces),
|
|
("Performance Benchmark", benchmark_performance),
|
|
("Integration Test", test_integration_with_existing_models)
|
|
]
|
|
|
|
results = {}
|
|
|
|
for test_name, test_func in tests:
|
|
try:
|
|
results[test_name] = test_func()
|
|
except Exception as e:
|
|
print(f"❌ {test_name} failed with exception: {e}")
|
|
results[test_name] = False
|
|
|
|
# Summary
|
|
print("\n" + "=" * 50)
|
|
print("TEST SUMMARY")
|
|
print("=" * 50)
|
|
|
|
passed = 0
|
|
total = len(tests)
|
|
|
|
for test_name, result in results.items():
|
|
status = "✅ PASS" if result else "❌ FAIL"
|
|
print(f"{test_name}: {status}")
|
|
if result:
|
|
passed += 1
|
|
|
|
print(f"\nOverall: {passed}/{total} tests passed")
|
|
|
|
if passed == total:
|
|
print("🎉 All NPU integration tests passed!")
|
|
else:
|
|
print("⚠️ Some tests failed. Check the output above for details.")
|
|
|
|
return passed == total
|
|
|
|
if __name__ == "__main__":
|
|
success = main()
|
|
sys.exit(0 if success else 1)
|
|
|