616 lines
20 KiB
Python
616 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Benchmark runner for PyTorch PINN - mirrors Rust benchmark structure.
|
|
|
|
This script benchmarks the PyTorch PINN implementation to compare with RustyTorch++.
|
|
|
|
Usage:
|
|
python benchmark_runner.py [--pytorch-pinn-path PATH] [--output OUTPUT]
|
|
|
|
The script expects a PyTorch PINN implementation with the following interface:
|
|
- Config class with physics/network/training parameters
|
|
- Mre1DPinnSolver class with:
|
|
- forward() method for network evaluation
|
|
- compute_pde_residual() method for PDE computation
|
|
- training_step() method for single training iteration
|
|
- train() method for full training loop
|
|
|
|
If no --pytorch-pinn-path is provided, a minimal reference implementation is used.
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
import gc
|
|
import argparse
|
|
import sys
|
|
from dataclasses import dataclass, asdict, field
|
|
from typing import List, Dict, Optional, Tuple
|
|
from pathlib import Path
|
|
|
|
# Optional imports with fallbacks
|
|
try:
|
|
import tracemalloc
|
|
HAS_TRACEMALLOC = True
|
|
except ImportError:
|
|
HAS_TRACEMALLOC = False
|
|
|
|
try:
|
|
import numpy as np
|
|
HAS_NUMPY = True
|
|
except ImportError:
|
|
HAS_NUMPY = False
|
|
print("Warning: numpy not available, some benchmarks will be skipped")
|
|
|
|
try:
|
|
import torch
|
|
import torch.nn as nn
|
|
HAS_TORCH = True
|
|
except ImportError:
|
|
HAS_TORCH = False
|
|
print("Warning: torch not available, PyTorch benchmarks will be skipped")
|
|
|
|
|
|
@dataclass
|
|
class BenchmarkResult:
|
|
"""Result of a single benchmark."""
|
|
name: str
|
|
n_points: int
|
|
mean_time_ms: float
|
|
std_time_ms: float
|
|
min_time_ms: float
|
|
max_time_ms: float
|
|
peak_memory_mb: float
|
|
iterations: int
|
|
throughput_elements_per_sec: float = 0.0
|
|
|
|
|
|
@dataclass
|
|
class BenchmarkSuite:
|
|
"""Collection of benchmark results."""
|
|
framework: str
|
|
device: str
|
|
version: str
|
|
benchmarks: List[BenchmarkResult] = field(default_factory=list)
|
|
|
|
|
|
def warmup(func, n_warmup: int = 10):
|
|
"""Run warmup iterations."""
|
|
for _ in range(n_warmup):
|
|
_ = func()
|
|
|
|
|
|
def measure_time(func, n_iterations: int = 100) -> Tuple[List[float], float]:
|
|
"""Measure execution time and peak memory."""
|
|
times = []
|
|
peak_memory = 0.0
|
|
|
|
if HAS_TRACEMALLOC:
|
|
tracemalloc.start()
|
|
|
|
if HAS_TORCH and torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
|
|
for _ in range(n_iterations):
|
|
start = time.perf_counter()
|
|
_ = func()
|
|
if HAS_TORCH and torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
times.append((time.perf_counter() - start) * 1000)
|
|
|
|
if HAS_TRACEMALLOC:
|
|
_, peak_memory = tracemalloc.get_traced_memory()
|
|
tracemalloc.stop()
|
|
peak_memory = peak_memory / 1024 / 1024 # Convert to MB
|
|
|
|
return times, peak_memory
|
|
|
|
|
|
def create_benchmark_result(
|
|
name: str,
|
|
n_points: int,
|
|
times: List[float],
|
|
peak_memory: float,
|
|
n_iterations: int
|
|
) -> BenchmarkResult:
|
|
"""Create a BenchmarkResult from timing data."""
|
|
import statistics
|
|
mean_time = statistics.mean(times)
|
|
return BenchmarkResult(
|
|
name=name,
|
|
n_points=n_points,
|
|
mean_time_ms=mean_time,
|
|
std_time_ms=statistics.stdev(times) if len(times) > 1 else 0.0,
|
|
min_time_ms=min(times),
|
|
max_time_ms=max(times),
|
|
peak_memory_mb=peak_memory,
|
|
iterations=n_iterations,
|
|
throughput_elements_per_sec=(n_points / mean_time * 1000) if mean_time > 0 else 0.0
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# MINIMAL PYTORCH PINN REFERENCE IMPLEMENTATION
|
|
# =============================================================================
|
|
|
|
if HAS_TORCH and HAS_NUMPY:
|
|
class MinimalConfig:
|
|
"""Minimal configuration matching Rust Config."""
|
|
def __init__(self, n_data: int = 200):
|
|
# Physics
|
|
self.rho = 1040.0
|
|
self.freq = 50.0
|
|
self.l = 0.1
|
|
self.u0 = 1e-6
|
|
self.g_prime_true = 3000.0
|
|
self.g_double_true = 1500.0
|
|
# Grid
|
|
self.n_data = n_data
|
|
self.n_pde = n_data
|
|
# Network
|
|
self.u_layers = 4
|
|
self.u_hidden = 64
|
|
self.u_ff_dim = 64
|
|
self.u_ff_scale = 10.0
|
|
# Training
|
|
self.lr = 1e-3
|
|
self.epochs = 50000
|
|
self.scheduler_patience = 500
|
|
self.scheduler_factor = 0.5
|
|
# Loss
|
|
self.data_weight = 1.0
|
|
self.pde_weight = 1e-6
|
|
|
|
def calculate_k(cfg: MinimalConfig) -> complex:
|
|
"""Calculate complex wave number."""
|
|
omega = 2.0 * np.pi * cfg.freq
|
|
G_complex = cfg.g_prime_true + 1j * cfg.g_double_true
|
|
k = np.sqrt(cfg.rho * omega**2 / G_complex)
|
|
if k.imag < 0:
|
|
k = -k
|
|
return k
|
|
|
|
def synthesize_displacement(cfg: MinimalConfig) -> Tuple[np.ndarray, np.ndarray, np.ndarray, complex]:
|
|
"""Generate synthetic displacement data."""
|
|
k = calculate_k(cfg)
|
|
x = np.linspace(0.0, cfg.l, cfg.n_data)
|
|
u_complex = cfg.u0 * np.exp(1j * k * x)
|
|
return x, u_complex.real, u_complex.imag, k
|
|
|
|
class LffnUNet1D(nn.Module):
|
|
"""Learnable Fourier Feature Network with MLP."""
|
|
def __init__(self, cfg: MinimalConfig):
|
|
super().__init__()
|
|
self.B_learnable = nn.Parameter(
|
|
torch.randn(1, cfg.u_ff_dim) * cfg.u_ff_scale
|
|
)
|
|
layers = []
|
|
dim = cfg.u_ff_dim * 2
|
|
for _ in range(cfg.u_layers):
|
|
layers.append(nn.Linear(dim, cfg.u_hidden))
|
|
layers.append(nn.Tanh())
|
|
dim = cfg.u_hidden
|
|
layers.append(nn.Linear(dim, 2))
|
|
self.net = nn.Sequential(*layers)
|
|
|
|
def forward(self, x_norm: torch.Tensor) -> torch.Tensor:
|
|
y = 2.0 * np.pi * x_norm @ self.B_learnable
|
|
sin_feat = torch.sin(y)
|
|
cos_feat = torch.cos(y)
|
|
feat = torch.cat([sin_feat, cos_feat], dim=-1)
|
|
return self.net(feat)
|
|
|
|
class MinimalPinnSolver:
|
|
"""Minimal PINN solver for benchmarking."""
|
|
def __init__(self, cfg: MinimalConfig, device: str = 'cpu'):
|
|
self.cfg = cfg
|
|
self.device = torch.device(device)
|
|
|
|
# Physics
|
|
self.omega = 2.0 * np.pi * cfg.freq
|
|
self.k_true = calculate_k(cfg)
|
|
|
|
# Generate data
|
|
x, u_r, u_i, _ = synthesize_displacement(cfg)
|
|
u_scale = np.max(np.sqrt(u_r**2 + u_i**2)) + 1e-16
|
|
self.u_scale = u_scale
|
|
|
|
# Normalize and convert to tensors
|
|
x_norm = torch.tensor(x / cfg.l, dtype=torch.float32, device=self.device).reshape(-1, 1)
|
|
u_data = torch.tensor(
|
|
np.stack([u_r / u_scale, u_i / u_scale], axis=1),
|
|
dtype=torch.float32, device=self.device
|
|
)
|
|
|
|
self.x_data = x_norm.requires_grad_(True)
|
|
self.u_data_target = u_data
|
|
|
|
# Network
|
|
self.u_net = LffnUNet1D(cfg).to(self.device)
|
|
|
|
# Optimizer
|
|
self.optimizer = torch.optim.Adam(self.u_net.parameters(), lr=cfg.lr)
|
|
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
|
self.optimizer, mode='min', factor=cfg.scheduler_factor,
|
|
patience=cfg.scheduler_patience
|
|
)
|
|
|
|
def forward(self) -> torch.Tensor:
|
|
"""Forward pass through network."""
|
|
return self.u_net(self.x_data)
|
|
|
|
def compute_data_loss(self, u_pred: torch.Tensor) -> torch.Tensor:
|
|
"""Compute MSE data loss."""
|
|
return torch.mean((u_pred - self.u_data_target)**2)
|
|
|
|
def compute_pde_residual(self) -> Tuple[float, float]:
|
|
"""Compute PDE residual using analytical derivatives."""
|
|
x = np.linspace(0.0, self.cfg.l, self.cfg.n_data)
|
|
u_complex = self.cfg.u0 * np.exp(1j * self.k_true * x)
|
|
|
|
# Analytical d²u/dx²
|
|
d2u = -(self.k_true**2) * u_complex
|
|
|
|
# Helmholtz: d²u/dx² + k²u = 0
|
|
k_sq = self.k_true**2
|
|
residual = d2u + k_sq * u_complex
|
|
|
|
mse_rr = np.mean(residual.real**2)
|
|
mse_ri = np.mean(residual.imag**2)
|
|
return mse_rr, mse_ri
|
|
|
|
def training_step(self) -> Tuple[float, float, float]:
|
|
"""Single training step."""
|
|
self.optimizer.zero_grad()
|
|
u_pred = self.forward()
|
|
loss_data = self.compute_data_loss(u_pred)
|
|
mse_rr, mse_ri = self.compute_pde_residual()
|
|
loss_pde = mse_rr + mse_ri
|
|
total_loss = self.cfg.data_weight * loss_data.item() + self.cfg.pde_weight * loss_pde
|
|
loss_data.backward()
|
|
self.optimizer.step()
|
|
self.scheduler.step(loss_data)
|
|
return loss_data.item(), loss_pde, total_loss
|
|
|
|
def train(self, epochs: int = None) -> float:
|
|
"""Full training loop."""
|
|
epochs = epochs or self.cfg.epochs
|
|
best_loss = float('inf')
|
|
for _ in range(epochs):
|
|
_, _, total_loss = self.training_step()
|
|
if total_loss < best_loss:
|
|
best_loss = total_loss
|
|
return best_loss
|
|
|
|
|
|
# =============================================================================
|
|
# BENCHMARK FUNCTIONS
|
|
# =============================================================================
|
|
|
|
def benchmark_forward_pass(suite: BenchmarkSuite, point_sizes: List[int] = [200, 1000, 10000], forced_device: str = "auto"):
|
|
"""Benchmark forward pass (network evaluation)."""
|
|
if not HAS_TORCH:
|
|
print("Skipping forward_pass benchmark: torch not available")
|
|
return
|
|
|
|
if forced_device == "auto":
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
elif forced_device == "cuda" and torch.cuda.is_available():
|
|
device = 'cuda'
|
|
else:
|
|
device = 'cpu'
|
|
|
|
for n_points in point_sizes:
|
|
print(f" forward_pass n={n_points}...", end=" ", flush=True)
|
|
|
|
cfg = MinimalConfig(n_data=n_points)
|
|
solver = MinimalPinnSolver(cfg, device=device)
|
|
|
|
warmup(solver.forward)
|
|
times, peak_memory = measure_time(solver.forward, n_iterations=100)
|
|
|
|
result = create_benchmark_result(
|
|
name="forward_pass/lffn_mlp",
|
|
n_points=n_points,
|
|
times=times,
|
|
peak_memory=peak_memory,
|
|
n_iterations=100
|
|
)
|
|
suite.benchmarks.append(result)
|
|
print(f"{result.mean_time_ms:.3f}ms")
|
|
|
|
del solver
|
|
gc.collect()
|
|
if HAS_TORCH and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def benchmark_pde_residual(suite: BenchmarkSuite, point_sizes: List[int] = [200, 1000, 10000], forced_device: str = "auto"):
|
|
"""Benchmark PDE residual computation."""
|
|
if not HAS_TORCH:
|
|
print("Skipping pde_residual benchmark: torch not available")
|
|
return
|
|
|
|
if forced_device == "auto":
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
elif forced_device == "cuda" and torch.cuda.is_available():
|
|
device = 'cuda'
|
|
else:
|
|
device = 'cpu'
|
|
|
|
for n_points in point_sizes:
|
|
print(f" pde_residual n={n_points}...", end=" ", flush=True)
|
|
|
|
cfg = MinimalConfig(n_data=n_points)
|
|
solver = MinimalPinnSolver(cfg, device=device)
|
|
|
|
warmup(solver.compute_pde_residual)
|
|
times, peak_memory = measure_time(solver.compute_pde_residual, n_iterations=100)
|
|
|
|
result = create_benchmark_result(
|
|
name="pde_residual/helmholtz",
|
|
n_points=n_points,
|
|
times=times,
|
|
peak_memory=peak_memory,
|
|
n_iterations=100
|
|
)
|
|
suite.benchmarks.append(result)
|
|
print(f"{result.mean_time_ms:.3f}ms")
|
|
|
|
del solver
|
|
gc.collect()
|
|
|
|
|
|
def benchmark_training_step(suite: BenchmarkSuite, point_sizes: List[int] = [200, 1000], forced_device: str = "auto"):
|
|
"""Benchmark single training step."""
|
|
if not HAS_TORCH:
|
|
print("Skipping training_step benchmark: torch not available")
|
|
return
|
|
|
|
if forced_device == "auto":
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
elif forced_device == "cuda" and torch.cuda.is_available():
|
|
device = 'cuda'
|
|
else:
|
|
device = 'cpu'
|
|
|
|
for n_points in point_sizes:
|
|
print(f" training_step n={n_points}...", end=" ", flush=True)
|
|
|
|
cfg = MinimalConfig(n_data=n_points)
|
|
solver = MinimalPinnSolver(cfg, device=device)
|
|
|
|
warmup(solver.training_step, n_warmup=5)
|
|
times, peak_memory = measure_time(solver.training_step, n_iterations=50)
|
|
|
|
result = create_benchmark_result(
|
|
name="training_step/single_step",
|
|
n_points=n_points,
|
|
times=times,
|
|
peak_memory=peak_memory,
|
|
n_iterations=50
|
|
)
|
|
suite.benchmarks.append(result)
|
|
print(f"{result.mean_time_ms:.3f}ms")
|
|
|
|
del solver
|
|
gc.collect()
|
|
if HAS_TORCH and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def benchmark_data_generation(suite: BenchmarkSuite, point_sizes: List[int] = [200, 1000, 10000, 100000]):
|
|
"""Benchmark data generation."""
|
|
if not HAS_NUMPY:
|
|
print("Skipping data_generation benchmark: numpy not available")
|
|
return
|
|
|
|
for n_points in point_sizes:
|
|
print(f" data_generation n={n_points}...", end=" ", flush=True)
|
|
|
|
cfg = MinimalConfig(n_data=n_points)
|
|
|
|
def gen_data():
|
|
return synthesize_displacement(cfg)
|
|
|
|
warmup(gen_data)
|
|
times, peak_memory = measure_time(gen_data, n_iterations=100)
|
|
|
|
result = create_benchmark_result(
|
|
name="data_generation/synthesize_displacement",
|
|
n_points=n_points,
|
|
times=times,
|
|
peak_memory=peak_memory,
|
|
n_iterations=100
|
|
)
|
|
suite.benchmarks.append(result)
|
|
print(f"{result.mean_time_ms:.3f}ms")
|
|
|
|
|
|
def benchmark_wave_number(suite: BenchmarkSuite):
|
|
"""Benchmark wave number calculation."""
|
|
if not HAS_NUMPY:
|
|
print("Skipping wave_number benchmark: numpy not available")
|
|
return
|
|
|
|
print(" calculate_k...", end=" ", flush=True)
|
|
cfg = MinimalConfig()
|
|
|
|
def calc_k():
|
|
return calculate_k(cfg)
|
|
|
|
warmup(calc_k)
|
|
times, peak_memory = measure_time(calc_k, n_iterations=1000)
|
|
|
|
result = create_benchmark_result(
|
|
name="wave_number/calculate_k",
|
|
n_points=1,
|
|
times=times,
|
|
peak_memory=peak_memory,
|
|
n_iterations=1000
|
|
)
|
|
suite.benchmarks.append(result)
|
|
print(f"{result.mean_time_ms:.4f}ms")
|
|
|
|
|
|
def benchmark_training_100_epochs(suite: BenchmarkSuite, point_sizes: List[int] = [200], forced_device: str = "auto"):
|
|
"""Benchmark 100 epochs of training."""
|
|
if not HAS_TORCH:
|
|
print("Skipping training_100_epochs benchmark: torch not available")
|
|
return
|
|
|
|
if forced_device == "auto":
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
elif forced_device == "cuda" and torch.cuda.is_available():
|
|
device = 'cuda'
|
|
else:
|
|
device = 'cpu'
|
|
|
|
for n_points in point_sizes:
|
|
print(f" training_100_epochs n={n_points}...", end=" ", flush=True)
|
|
|
|
def run_training():
|
|
cfg = MinimalConfig(n_data=n_points)
|
|
solver = MinimalPinnSolver(cfg, device=device)
|
|
return solver.train(epochs=100)
|
|
|
|
# Only 5 iterations since this is slow
|
|
times, peak_memory = measure_time(run_training, n_iterations=5)
|
|
|
|
result = create_benchmark_result(
|
|
name="training_100_epochs/train_100",
|
|
n_points=n_points,
|
|
times=times,
|
|
peak_memory=peak_memory,
|
|
n_iterations=5
|
|
)
|
|
suite.benchmarks.append(result)
|
|
print(f"{result.mean_time_ms:.1f}ms")
|
|
|
|
gc.collect()
|
|
if HAS_TORCH and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def run_all_benchmarks(point_sizes: Optional[List[int]] = None, forced_device: str = "auto") -> BenchmarkSuite:
|
|
"""Run complete benchmark suite.
|
|
|
|
Args:
|
|
point_sizes: List of point sizes to benchmark
|
|
forced_device: Device to use - "cpu", "cuda", or "auto" (detect)
|
|
"""
|
|
point_sizes = point_sizes or [200, 1000, 10000]
|
|
|
|
# Determine device based on forced_device argument
|
|
if forced_device == "auto":
|
|
if HAS_TORCH and torch.cuda.is_available():
|
|
device = f"cuda:{torch.cuda.get_device_name(0)}"
|
|
else:
|
|
device = "cpu"
|
|
elif forced_device == "cuda":
|
|
if HAS_TORCH and torch.cuda.is_available():
|
|
device = f"cuda:{torch.cuda.get_device_name(0)}"
|
|
else:
|
|
print("Warning: CUDA requested but not available, falling back to CPU")
|
|
device = "cpu"
|
|
else:
|
|
device = "cpu"
|
|
|
|
version = "unknown"
|
|
if HAS_TORCH:
|
|
version = torch.__version__
|
|
|
|
suite = BenchmarkSuite(
|
|
framework="pytorch",
|
|
device=device,
|
|
version=version
|
|
)
|
|
|
|
print(f"\n=== PyTorch PINN Benchmarks ===")
|
|
print(f"Device: {device}")
|
|
print(f"Version: {version}")
|
|
print()
|
|
|
|
print("Running benchmarks:")
|
|
benchmark_wave_number(suite)
|
|
benchmark_data_generation(suite, point_sizes + [100000])
|
|
benchmark_forward_pass(suite, point_sizes, forced_device)
|
|
benchmark_pde_residual(suite, point_sizes, forced_device)
|
|
benchmark_training_step(suite, [200, 1000], forced_device)
|
|
benchmark_training_100_epochs(suite, [200], forced_device)
|
|
|
|
return suite
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="PyTorch PINN Benchmark Runner")
|
|
parser.add_argument(
|
|
"--pytorch-pinn-path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to PyTorch PINN implementation (uses minimal reference if not provided)"
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=str,
|
|
default="python_results.json",
|
|
help="Output JSON file for results"
|
|
)
|
|
parser.add_argument(
|
|
"--point-sizes",
|
|
type=str,
|
|
default="200,1000,10000",
|
|
help="Comma-separated list of point sizes to benchmark"
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
type=str,
|
|
choices=["cpu", "cuda", "auto"],
|
|
default="auto",
|
|
help="Device to use: cpu, cuda, or auto (detect)"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
point_sizes = [int(x) for x in args.point_sizes.split(",")]
|
|
forced_device = args.device
|
|
|
|
if args.pytorch_pinn_path:
|
|
print(f"Loading custom PyTorch PINN from: {args.pytorch_pinn_path}")
|
|
# Import the reference implementation if available
|
|
import importlib.util
|
|
spec = importlib.util.spec_from_file_location("pytorch_pinn", args.pytorch_pinn_path)
|
|
if spec and spec.loader:
|
|
pytorch_pinn = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(pytorch_pinn)
|
|
print(f"Loaded custom PINN from: {args.pytorch_pinn_path}")
|
|
# Use the loaded implementation for benchmarks
|
|
global MinimalConfig, MinimalPinnSolver, LffnUNet1D
|
|
MinimalConfig = pytorch_pinn.Config
|
|
# Note: Would need adapter for Mre1DFixedGPinn -> MinimalPinnSolver interface
|
|
else:
|
|
print("Failed to load custom PINN, using minimal reference")
|
|
|
|
suite = run_all_benchmarks(point_sizes, forced_device)
|
|
|
|
# Save results
|
|
output_path = Path(args.output)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
with open(output_path, "w") as f:
|
|
json.dump(asdict(suite), f, indent=2)
|
|
|
|
print(f"\nResults saved to: {output_path}")
|
|
|
|
# Print summary table
|
|
print("\n" + "=" * 70)
|
|
print("BENCHMARK SUMMARY")
|
|
print("=" * 70)
|
|
print(f"{'Benchmark':<40} {'Points':>8} {'Mean (ms)':>12} {'Std':>10}")
|
|
print("-" * 70)
|
|
for result in suite.benchmarks:
|
|
print(f"{result.name:<40} {result.n_points:>8} {result.mean_time_ms:>12.3f} {result.std_time_ms:>10.3f}")
|
|
print("=" * 70)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|