initial commit

benchmark.py

@@ -0,0 +1,176 @@
+import logging
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from config import config
+from src.amd.optimizations import AMDOptimizer
+from src.models.intraday_gnn import IntradayGNN
+from src.utils.memory_manager import MemoryManager
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    handlers=[
+        logging.FileHandler("benchmark_r9700.log"),
+        logging.StreamHandler(),
+    ],
+)
+logger = logging.getLogger(__name__)
+
+
+def benchmark_model():
+    """Benchmark the GNN model on AMD Radeon R9700 AI Pro"""
+    # Initialize memory manager
+    memory_manager = MemoryManager()
+    logger.info(memory_manager.get_memory_stats())
+
+    # Initialize AMD optimizer
+    amd_optimizer = AMDOptimizer()
+
+    # Create a sample model
+    num_features = len(config.INTRADAY_FEATURES) + 5  # +5 for price features
+    model = IntradayGNN(num_features, config.SEQUENCE_LENGTH)
+
+    # Optimize model for AMD GPU
+    model = amd_optimizer.optimize_model(model)
+
+    # Create sample data
+    batch_size = config.BATCH_SIZE
+    sequence_length = config.SEQUENCE_LENGTH
+    num_stocks = 50  # Number of stocks in the graph
+
+    # Create random data
+    x = torch.randn(num_stocks, sequence_length, num_features).to(config.DEVICE)
+
+    # Create random edges
+    num_edges = 200
+    edge_index = torch.randint(0, num_stocks, (2, num_edges)).to(config.DEVICE)
+    edge_attr = torch.randn(num_edges, 1).to(config.DEVICE)
+
+    # Create target
+    y = torch.randn(num_stocks, 1).to(config.DEVICE)
+
+    # Warm-up
+    logger.info("Warming up...")
+    for _ in range(10):
+        with torch.no_grad():
+            _ = model((x, edge_index, edge_attr))
+
+    # Benchmark inference
+    logger.info("Benchmarking inference...")
+    start_time = time.time()
+    num_runs = 100
+
+    for _ in range(num_runs):
+        with torch.no_grad():
+            _ = model((x, edge_index, edge_attr))
+
+    inference_time = (time.time() - start_time) / num_runs
+    logger.info(f"Average inference time: {inference_time:.6f} seconds")
+
+    # Benchmark training
+    logger.info("Benchmarking training...")
+    model.train()
+    optimizer = torch.optim.Adam(model.parameters(), lr=config.LEARNING_RATE)
+    criterion = nn.MSELoss()
+
+    start_time = time.time()
+
+    for _ in range(num_runs):
+        optimizer.zero_grad()
+        out = model((x, edge_index, edge_attr))
+        loss = criterion(out, y)
+        loss.backward()
+        optimizer.step()
+
+    training_time = (time.time() - start_time) / num_runs
+    logger.info(f"Average training time: {training_time:.6f} seconds")
+
+    # Memory usage
+    memory_info = memory_manager.check_memory()
+    logger.info(
+        f"GPU Memory Usage: {memory_info['allocated'] / 1024**3:.2f}GB / {memory_info['total'] / 1024**3:.2f}GB"
+    )
+
+    # Throughput
+    logger.info(f"Inference throughput: {1 / inference_time:.2f} samples/second")
+    logger.info(f"Training throughput: {1 / training_time:.2f} samples/second")
+
+    # Detailed benchmark with different batch sizes
+    logger.info("\nDetailed benchmark with different configurations:")
+
+    batch_sizes = [32, 64, 128, 256]
+    sequence_lengths = [30, 60, 120]
+
+    for batch_size in batch_sizes:
+        for seq_len in sequence_lengths:
+            # Create data for this configuration
+            x = torch.randn(num_stocks, seq_len, num_features).to(config.DEVICE)
+            edge_index = torch.randint(0, num_stocks, (2, num_edges)).to(config.DEVICE)
+            edge_attr = torch.randn(num_edges, 1).to(config.DEVICE)
+            y = torch.randn(num_stocks, 1).to(config.DEVICE)
+
+            # Benchmark inference
+            start_time = time.time()
+            for _ in range(10):  # Fewer runs for detailed benchmark
+                with torch.no_grad():
+                    _ = model((x, edge_index, edge_attr))
+            inf_time = (time.time() - start_time) / 10
+
+            # Benchmark training
+            start_time = time.time()
+            for _ in range(10):
+                optimizer.zero_grad()
+                out = model((x, edge_index, edge_attr))
+                loss = criterion(out, y)
+                loss.backward()
+                optimizer.step()
+            train_time = (time.time() - start_time) / 10
+
+            logger.info(
+                f"Batch: {batch_size}, Seq Len: {seq_len}, "
+                f"Inf Time: {inf_time:.6f}s, Train Time: {train_time:.6f}s, "
+                f"Inf Tput: {1 / inf_time:.2f} samples/s, Train Tput: {1 / train_time:.2f} samples/s"
+            )
+
+    # Memory benchmark
+    logger.info("\nMemory benchmark:")
+
+    # Test different model sizes
+    hidden_channels_list = [64, 128, 256, 512]
+
+    for hidden_channels in hidden_channels_list:
+        # Create a model with this configuration
+        model = IntradayGNN(num_features, config.SEQUENCE_LENGTH)
+        model.feature_processor = nn.Sequential(
+            nn.Linear(num_features, hidden_channels),
+            nn.SiLU(),
+            nn.Linear(hidden_channels, hidden_channels),
+            nn.LayerNorm(hidden_channels),
+        )
+        model.linear = nn.Linear(hidden_channels, 1)
+
+        # Optimize model
+        model = amd_optimizer.optimize_model(model)
+
+        # Estimate memory usage
+        estimated_memory = memory_manager.estimate_model_memory(model)
+        logger.info(
+            f"Hidden Channels: {hidden_channels}, Estimated Memory: {estimated_memory / 1024**3:.2f}GB"
+        )
+
+        # Clean up
+        del model
+        memory_manager.empty_cache()
+
+    # Final memory stats
+    logger.info("\nFinal Memory Stats:")
+    logger.info(memory_manager.get_memory_stats())
+
+
+if __name__ == "__main__":
+    benchmark_model()