297 lines
8.6 KiB
Python
297 lines
8.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
PyTorch Benchmark for PINN MRE Helmholtz solver.
|
|
|
|
This script benchmarks forward pass and training step performance on CPU and GPU
|
|
for direct comparison with the Rust implementation.
|
|
|
|
Run with: python benchmark_pytorch.py
|
|
"""
|
|
|
|
import time
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from dataclasses import dataclass
|
|
import statistics
|
|
|
|
# Disable gradient computation for forward-only benchmarks
|
|
torch.set_grad_enabled(True)
|
|
|
|
@dataclass
|
|
class Config:
|
|
"""Configuration matching the Rust implementation."""
|
|
RHO: float = 1040.0
|
|
FREQ: float = 50.0
|
|
L: float = 0.1
|
|
U0: float = 1e-6
|
|
G_PRIME_TRUE: float = 3000.0
|
|
G_DOUBLE_TRUE: float = 1500.0
|
|
N_DATA: int = 200
|
|
N_PDE: int = 200
|
|
U_LAYERS: int = 4
|
|
U_HIDDEN: int = 64
|
|
U_FF_DIM: int = 64
|
|
U_FF_SCALE: float = 10.0
|
|
LR: float = 1e-3
|
|
EPOCHS: int = 100
|
|
|
|
|
|
def calculate_k(cfg: Config) -> complex:
|
|
"""Calculates the complex wave number k."""
|
|
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 np.imag(k) < 0:
|
|
k = -k
|
|
return k
|
|
|
|
|
|
def synthesize_displacement(cfg: Config, device: torch.device) -> tuple:
|
|
"""Generates displacement data u(x)."""
|
|
k = calculate_k(cfg)
|
|
x = np.linspace(0.0, cfg.L, cfg.N_DATA)
|
|
u_complex = cfg.U0 * np.exp(1j * k * x)
|
|
|
|
x_norm = torch.from_numpy(x / cfg.L).float().reshape(-1, 1).to(device)
|
|
u_data = torch.from_numpy(
|
|
np.stack([np.real(u_complex), np.imag(u_complex)], axis=1)
|
|
).float().to(device)
|
|
|
|
return x_norm, u_data, k
|
|
|
|
|
|
class LffnUNet1D(nn.Module):
|
|
"""Learnable Fourier Feature Network (LFFN-MLP) - matches Rust implementation."""
|
|
def __init__(self, cfg: Config):
|
|
super().__init__()
|
|
self.cfg = cfg
|
|
|
|
# Learnable Fourier features
|
|
B_init = torch.randn(1, cfg.U_FF_DIM) * cfg.U_FF_SCALE
|
|
self.B_learnable = nn.Parameter(B_init, requires_grad=True)
|
|
|
|
# Build MLP
|
|
in_dim = cfg.U_FF_DIM * 2
|
|
layers = []
|
|
dim = in_dim
|
|
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)
|
|
feat = torch.cat([torch.sin(y), torch.cos(y)], dim=-1)
|
|
return self.net(feat)
|
|
|
|
|
|
def benchmark_forward_pass(n_points: int, device: torch.device, n_warmup: int = 10, n_iter: int = 50) -> dict:
|
|
"""Benchmark forward pass only."""
|
|
cfg = Config(N_DATA=n_points, N_PDE=n_points)
|
|
|
|
# Setup
|
|
x_norm, u_data, k = synthesize_displacement(cfg, device)
|
|
x_norm.requires_grad_(True)
|
|
|
|
model = LffnUNet1D(cfg).to(device)
|
|
model.eval()
|
|
|
|
# Warmup
|
|
for _ in range(n_warmup):
|
|
with torch.no_grad():
|
|
_ = model(x_norm)
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
# Benchmark
|
|
times = []
|
|
for _ in range(n_iter):
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
start = time.perf_counter()
|
|
with torch.no_grad():
|
|
output = model(x_norm)
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
end = time.perf_counter()
|
|
times.append((end - start) * 1000) # Convert to ms
|
|
|
|
return {
|
|
'mean_ms': statistics.mean(times),
|
|
'std_ms': statistics.stdev(times) if len(times) > 1 else 0,
|
|
'min_ms': min(times),
|
|
'max_ms': max(times),
|
|
'throughput_kelem_s': (n_points / statistics.mean(times)),
|
|
}
|
|
|
|
|
|
def benchmark_training_step(n_points: int, device: torch.device, n_warmup: int = 5, n_iter: int = 30) -> dict:
|
|
"""Benchmark a single training step (forward + backward + optimizer step)."""
|
|
cfg = Config(N_DATA=n_points, N_PDE=n_points)
|
|
|
|
# Setup
|
|
x_norm, u_data, k = synthesize_displacement(cfg, device)
|
|
x_norm.requires_grad_(True)
|
|
|
|
# Normalize data (matching Rust)
|
|
u_scale = float(torch.abs(u_data).max().item() + 1e-16)
|
|
u_data_norm = u_data / u_scale
|
|
|
|
model = LffnUNet1D(cfg).to(device)
|
|
model.train()
|
|
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=cfg.LR)
|
|
mse_loss = nn.MSELoss()
|
|
|
|
# Warmup
|
|
for _ in range(n_warmup):
|
|
optimizer.zero_grad()
|
|
output = model(x_norm)
|
|
loss = mse_loss(output, u_data_norm)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
# Benchmark
|
|
times = []
|
|
for _ in range(n_iter):
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
start = time.perf_counter()
|
|
|
|
optimizer.zero_grad()
|
|
output = model(x_norm)
|
|
loss = mse_loss(output, u_data_norm)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
end = time.perf_counter()
|
|
times.append((end - start) * 1000) # Convert to ms
|
|
|
|
return {
|
|
'mean_ms': statistics.mean(times),
|
|
'std_ms': statistics.stdev(times) if len(times) > 1 else 0,
|
|
'min_ms': min(times),
|
|
'max_ms': max(times),
|
|
'throughput_kelem_s': (n_points / statistics.mean(times)),
|
|
}
|
|
|
|
|
|
def benchmark_training_100_epochs(n_points: int, device: torch.device, n_iter: int = 5) -> dict:
|
|
"""Benchmark 100 epochs of training."""
|
|
cfg = Config(N_DATA=n_points, N_PDE=n_points, EPOCHS=100)
|
|
|
|
times = []
|
|
for _ in range(n_iter):
|
|
# Fresh setup for each run
|
|
x_norm, u_data, k = synthesize_displacement(cfg, device)
|
|
x_norm.requires_grad_(True)
|
|
|
|
u_scale = float(torch.abs(u_data).max().item() + 1e-16)
|
|
u_data_norm = u_data / u_scale
|
|
|
|
model = LffnUNet1D(cfg).to(device)
|
|
model.train()
|
|
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=cfg.LR)
|
|
mse_loss = nn.MSELoss()
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
start = time.perf_counter()
|
|
|
|
for epoch in range(100):
|
|
optimizer.zero_grad()
|
|
output = model(x_norm)
|
|
loss = mse_loss(output, u_data_norm)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
end = time.perf_counter()
|
|
times.append((end - start) * 1000) # Convert to ms
|
|
|
|
return {
|
|
'mean_ms': statistics.mean(times),
|
|
'std_ms': statistics.stdev(times) if len(times) > 1 else 0,
|
|
'min_ms': min(times),
|
|
'max_ms': max(times),
|
|
'epochs_per_sec': 100 / (statistics.mean(times) / 1000),
|
|
}
|
|
|
|
|
|
def print_results(name: str, results: dict):
|
|
"""Print benchmark results in a formatted way."""
|
|
print(f" {name}:")
|
|
print(f" Time: {results['mean_ms']:.3f} ms (+/- {results['std_ms']:.3f} ms)")
|
|
print(f" Min: {results['min_ms']:.3f} ms, Max: {results['max_ms']:.3f} ms")
|
|
if 'throughput_kelem_s' in results:
|
|
print(f" Throughput: {results['throughput_kelem_s']:.2f} Kelem/s")
|
|
if 'epochs_per_sec' in results:
|
|
print(f" Epochs/sec: {results['epochs_per_sec']:.2f}")
|
|
|
|
|
|
def main():
|
|
print("=" * 70)
|
|
print("PyTorch PINN MRE Helmholtz Benchmark")
|
|
print("=" * 70)
|
|
|
|
# Check available devices
|
|
devices = [('cpu', torch.device('cpu'))]
|
|
if torch.cuda.is_available():
|
|
devices.append(('cuda', torch.device('cuda')))
|
|
print(f"CUDA available: {torch.cuda.get_device_name(0)}")
|
|
else:
|
|
print("CUDA not available, running CPU benchmarks only")
|
|
|
|
print(f"PyTorch version: {torch.__version__}")
|
|
print()
|
|
|
|
batch_sizes = [200, 1000, 10000]
|
|
|
|
for device_name, device in devices:
|
|
print(f"\n{'=' * 70}")
|
|
print(f"Device: {device_name.upper()}")
|
|
print(f"{'=' * 70}")
|
|
|
|
# Forward pass benchmarks
|
|
print("\n--- Forward Pass (inference only) ---")
|
|
for n_points in batch_sizes:
|
|
results = benchmark_forward_pass(n_points, device)
|
|
print_results(f"n_points={n_points}", results)
|
|
|
|
# Training step benchmarks
|
|
print("\n--- Single Training Step (forward + backward + optimizer) ---")
|
|
for n_points in batch_sizes[:2]: # Only 200 and 1000 for training
|
|
results = benchmark_training_step(n_points, device)
|
|
print_results(f"n_points={n_points}", results)
|
|
|
|
# 100 epochs benchmark
|
|
print("\n--- 100 Epochs Training ---")
|
|
results = benchmark_training_100_epochs(200, device)
|
|
print_results("n_points=200, epochs=100", results)
|
|
|
|
print("\n" + "=" * 70)
|
|
print("Benchmark complete!")
|
|
print("=" * 70)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|