Files
rustytorch/demos/rtx-bioheat/python/run_comparison.py
T
2026-03-04 00:08:42 +00:00

377 lines
13 KiB
Python

#!/usr/bin/env python3
"""
Benchmark Comparison Script for Bioheat PINN
Compares RustyTorch++ (Rust) vs PyTorch (Python) performance
for solving the Pennes Bioheat Equation using PINNs.
Usage:
python run_comparison.py [--steps 1000] [--output comparison_report.html]
"""
import argparse
import json
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, Optional
# Import the PyTorch benchmark
from pytorch_bioheat_reference import benchmark_pytorch
def run_rustytorch_benchmark(
num_steps: int = 1000,
binary_path: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
"""
Run RustyTorch++ benchmark via CLI.
Note: This requires a compiled benchmark binary.
If not available, returns None.
"""
# Look for benchmark binary
if binary_path is None:
possible_paths = [
"../../target/release/bioheat_benchmark",
"../../target/debug/bioheat_benchmark",
"../target/release/bioheat_benchmark",
]
for path in possible_paths:
if os.path.exists(path):
binary_path = path
break
if binary_path is None or not os.path.exists(binary_path):
print("RustyTorch++ benchmark binary not found.")
print("To compile: cargo build --release -p rtx-bioheat --bin bioheat_benchmark")
return None
try:
result = subprocess.run(
[binary_path, "--steps", str(num_steps), "--json"],
capture_output=True,
text=True,
timeout=600,
)
if result.returncode != 0:
print(f"RustyTorch++ benchmark failed: {result.stderr}")
return None
return json.loads(result.stdout)
except Exception as e:
print(f"Error running RustyTorch++ benchmark: {e}")
return None
def generate_html_report(
pytorch_results: Dict[str, Any],
rustytorch_results: Optional[Dict[str, Any]],
output_path: str,
):
"""Generate HTML comparison report."""
# Calculate speedups if RustyTorch results available
if rustytorch_results:
training_speedup = pytorch_results['training_time_sec'] / rustytorch_results['training_time_sec']
inference_speedup = pytorch_results['inference_mean_ms'] / rustytorch_results['inference_mean_ms']
else:
training_speedup = None
inference_speedup = None
html = f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Bioheat PINN Benchmark Comparison</title>
<style>
body {{
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
max-width: 1200px;
margin: 0 auto;
padding: 40px 20px;
background: #0a0a0f;
color: #e0e0e0;
}}
h1 {{
color: #4CAF50;
text-align: center;
margin-bottom: 10px;
}}
.subtitle {{
text-align: center;
color: #888;
margin-bottom: 40px;
}}
.section {{
background: #1a1a2e;
border-radius: 12px;
padding: 25px;
margin-bottom: 25px;
}}
h2 {{
color: #4CAF50;
margin-top: 0;
border-bottom: 1px solid #333;
padding-bottom: 10px;
}}
.comparison-table {{
width: 100%;
border-collapse: collapse;
margin-top: 15px;
}}
.comparison-table th, .comparison-table td {{
padding: 12px 15px;
text-align: left;
border-bottom: 1px solid #333;
}}
.comparison-table th {{
background: #252540;
color: #4CAF50;
}}
.comparison-table tr:hover {{
background: #252540;
}}
.speedup {{
color: #4CAF50;
font-weight: bold;
}}
.metric-card {{
display: inline-block;
background: #252540;
border-radius: 8px;
padding: 20px;
margin: 10px;
min-width: 200px;
text-align: center;
}}
.metric-value {{
font-size: 32px;
font-weight: bold;
color: #4CAF50;
}}
.metric-label {{
color: #888;
margin-top: 5px;
}}
.chart-container {{
margin-top: 20px;
height: 300px;
}}
.bar {{
display: inline-block;
width: 80px;
margin: 0 20px;
vertical-align: bottom;
}}
.bar-pytorch {{
background: linear-gradient(to top, #EE4C2C, #FF6F61);
}}
.bar-rustytorch {{
background: linear-gradient(to top, #4CAF50, #8BC34A);
}}
.bar-label {{
text-align: center;
margin-top: 10px;
color: #888;
}}
.footer {{
text-align: center;
color: #666;
margin-top: 40px;
padding-top: 20px;
border-top: 1px solid #333;
}}
.equation {{
font-family: 'Times New Roman', serif;
font-size: 18px;
text-align: center;
padding: 20px;
background: #252540;
border-radius: 8px;
margin: 15px 0;
}}
</style>
</head>
<body>
<h1>Bioheat PINN Benchmark Comparison</h1>
<p class="subtitle">RustyTorch++ vs PyTorch - Pennes Bioheat Equation</p>
<div class="section">
<h2>Physics: Pennes Bioheat Equation</h2>
<div class="equation">
&rho;c(&part;T/&part;t) = k&nabla;&sup2;T + &omega;<sub>b</sub>&rho;<sub>b</sub>c<sub>b</sub>(T<sub>a</sub> - T) + Q<sub>m</sub> + Q<sub>s</sub>
</div>
<p>The Pennes bioheat equation models heat transfer in biological tissue for thermal ablation planning.</p>
</div>
<div class="section">
<h2>Summary Metrics</h2>
<div style="text-align: center;">
<div class="metric-card">
<div class="metric-value">{pytorch_results['steps_per_second']:.0f}</div>
<div class="metric-label">PyTorch Steps/sec</div>
</div>
{"<div class='metric-card'><div class='metric-value'>" + f"{rustytorch_results['steps_per_second']:.0f}" + "</div><div class='metric-label'>RustyTorch++ Steps/sec</div></div>" if rustytorch_results else ""}
{f"<div class='metric-card'><div class='metric-value speedup'>{training_speedup:.1f}x</div><div class='metric-label'>Training Speedup</div></div>" if training_speedup else ""}
{f"<div class='metric-card'><div class='metric-value speedup'>{inference_speedup:.1f}x</div><div class='metric-label'>Inference Speedup</div></div>" if inference_speedup else ""}
</div>
</div>
<div class="section">
<h2>Detailed Comparison</h2>
<table class="comparison-table">
<thead>
<tr>
<th>Metric</th>
<th>PyTorch</th>
{"<th>RustyTorch++</th>" if rustytorch_results else ""}
{"<th>Speedup</th>" if rustytorch_results else ""}
</tr>
</thead>
<tbody>
<tr>
<td>Device</td>
<td>{pytorch_results['device']}</td>
{f"<td>{rustytorch_results['device']}</td>" if rustytorch_results else ""}
{f"<td>-</td>" if rustytorch_results else ""}
</tr>
<tr>
<td>Model Parameters</td>
<td>{pytorch_results['num_parameters']:,}</td>
{f"<td>{rustytorch_results['num_parameters']:,}</td>" if rustytorch_results else ""}
{f"<td>-</td>" if rustytorch_results else ""}
</tr>
<tr>
<td>Training Steps</td>
<td>{pytorch_results['num_steps']}</td>
{f"<td>{rustytorch_results['num_steps']}</td>" if rustytorch_results else ""}
{f"<td>-</td>" if rustytorch_results else ""}
</tr>
<tr>
<td>Training Time (sec)</td>
<td>{pytorch_results['training_time_sec']:.2f}</td>
{f"<td>{rustytorch_results['training_time_sec']:.2f}</td>" if rustytorch_results else ""}
{f"<td class='speedup'>{training_speedup:.2f}x</td>" if training_speedup else ""}
</tr>
<tr>
<td>Steps/Second</td>
<td>{pytorch_results['steps_per_second']:.1f}</td>
{f"<td>{rustytorch_results['steps_per_second']:.1f}</td>" if rustytorch_results else ""}
{f"<td class='speedup'>{training_speedup:.2f}x</td>" if training_speedup else ""}
</tr>
<tr>
<td>Inference Time (32^3 grid)</td>
<td>{pytorch_results['inference_mean_ms']:.2f} ms</td>
{f"<td>{rustytorch_results['inference_mean_ms']:.2f} ms</td>" if rustytorch_results else ""}
{f"<td class='speedup'>{inference_speedup:.2f}x</td>" if inference_speedup else ""}
</tr>
<tr>
<td>Peak GPU Memory</td>
<td>{pytorch_results['peak_memory_mb']:.1f} MB</td>
{f"<td>{rustytorch_results['peak_memory_mb']:.1f} MB</td>" if rustytorch_results else ""}
{f"<td>-</td>" if rustytorch_results else ""}
</tr>
<tr>
<td>Final Loss</td>
<td>{pytorch_results['final_loss']:.4e}</td>
{f"<td>{rustytorch_results['final_loss']:.4e}</td>" if rustytorch_results else ""}
{f"<td>-</td>" if rustytorch_results else ""}
</tr>
</tbody>
</table>
</div>
<div class="section">
<h2>Configuration</h2>
<ul>
<li>Architecture: 4D PINN with Fourier Features + MLP</li>
<li>Input: (x, y, z, t) spatial-temporal coordinates</li>
<li>Output: Temperature T(x,y,z,t)</li>
<li>Fourier Features: 64</li>
<li>Hidden Layers: [128, 128, 128, 64]</li>
<li>Activation: Tanh</li>
<li>Collocation Points: 4096</li>
<li>Tissue Type: Liver</li>
<li>Probe Power: 20W</li>
</ul>
</div>
<div class="footer">
<p>Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
<p>RustyTorch++ Bioheat PINN Benchmark</p>
</div>
</body>
</html>"""
with open(output_path, 'w') as f:
f.write(html)
print(f"\nReport generated: {output_path}")
def main():
parser = argparse.ArgumentParser(
description='Benchmark Comparison: RustyTorch++ vs PyTorch'
)
parser.add_argument(
'--steps', type=int, default=1000,
help='Number of training steps (default: 1000)'
)
parser.add_argument(
'--output', type=str, default='comparison_report.html',
help='Output HTML report path (default: comparison_report.html)'
)
parser.add_argument(
'--pytorch-only', action='store_true',
help='Only run PyTorch benchmark'
)
parser.add_argument(
'--device', type=str, default='cuda',
choices=['cuda', 'cpu'],
help='Device for PyTorch (default: cuda)'
)
args = parser.parse_args()
print("=" * 60)
print("Bioheat PINN Benchmark Comparison")
print("RustyTorch++ vs PyTorch")
print("=" * 60)
# Run PyTorch benchmark
print("\n[1/2] Running PyTorch benchmark...")
pytorch_results = benchmark_pytorch(
device=args.device,
num_steps=args.steps,
)
# Run RustyTorch++ benchmark (if available)
rustytorch_results = None
if not args.pytorch_only:
print("\n[2/2] Running RustyTorch++ benchmark...")
rustytorch_results = run_rustytorch_benchmark(num_steps=args.steps)
# Generate report
print("\nGenerating comparison report...")
generate_html_report(pytorch_results, rustytorch_results, args.output)
# Print summary
print("\n" + "=" * 60)
print("Summary")
print("=" * 60)
print(f"PyTorch training: {pytorch_results['steps_per_second']:.1f} steps/sec")
if rustytorch_results:
print(f"RustyTorch++ training: {rustytorch_results['steps_per_second']:.1f} steps/sec")
speedup = pytorch_results['training_time_sec'] / rustytorch_results['training_time_sec']
print(f"Training speedup: {speedup:.2f}x")
print("=" * 60)
if __name__ == '__main__':
main()