724 lines
22 KiB
Python
724 lines
22 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
PyTorch Reference Implementation for Pennes Bioheat PINN
|
|
|
|
This script provides a reference implementation of the Physics-Informed Neural Network
|
|
for solving the Pennes Bioheat Equation, used to benchmark against RustyTorch++.
|
|
|
|
Pennes Bioheat Equation:
|
|
rho*c * dT/dt = k * laplacian(T) + omega_b*rho_b*c_b*(T_a - T) + Q_m + Q_s
|
|
|
|
Where:
|
|
- T: Temperature field T(x,y,z,t)
|
|
- rho, c: Tissue density and specific heat
|
|
- k: Thermal conductivity
|
|
- omega_b: Blood perfusion rate
|
|
- T_a: Arterial blood temperature (37C)
|
|
- Q_m: Metabolic heat generation
|
|
- Q_s: External heat source (ablation probe)
|
|
|
|
Usage:
|
|
python pytorch_bioheat_reference.py [--device cuda|cpu] [--steps 1000]
|
|
"""
|
|
|
|
import argparse
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Tuple, Optional, Dict, Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
# =============================================================================
|
|
# Tissue Properties
|
|
# =============================================================================
|
|
|
|
@dataclass
|
|
class TissueProperties:
|
|
"""Physical properties for biological tissue."""
|
|
name: str
|
|
density: float # kg/m^3
|
|
specific_heat: float # J/(kg*K)
|
|
conductivity: float # W/(m*K)
|
|
perfusion: float # 1/s (blood perfusion rate)
|
|
metabolic_heat: float # W/m^3
|
|
|
|
@classmethod
|
|
def liver(cls) -> 'TissueProperties':
|
|
return cls(
|
|
name="Liver",
|
|
density=1060.0,
|
|
specific_heat=3600.0,
|
|
conductivity=0.512,
|
|
perfusion=0.0064,
|
|
metabolic_heat=420.0
|
|
)
|
|
|
|
@classmethod
|
|
def kidney(cls) -> 'TissueProperties':
|
|
return cls(
|
|
name="Kidney",
|
|
density=1050.0,
|
|
specific_heat=3890.0,
|
|
conductivity=0.544,
|
|
perfusion=0.0083,
|
|
metabolic_heat=380.0
|
|
)
|
|
|
|
@classmethod
|
|
def tumor(cls) -> 'TissueProperties':
|
|
return cls(
|
|
name="Tumor",
|
|
density=1040.0,
|
|
specific_heat=3800.0,
|
|
conductivity=0.55,
|
|
perfusion=0.002, # Lower perfusion = hotter
|
|
metabolic_heat=500.0
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class BloodProperties:
|
|
"""Physical properties for blood."""
|
|
density: float = 1060.0 # kg/m^3
|
|
specific_heat: float = 3770.0 # J/(kg*K)
|
|
arterial_temp: float = 37.0 # Celsius
|
|
|
|
|
|
# =============================================================================
|
|
# PINN Network Architecture
|
|
# =============================================================================
|
|
|
|
class FourierFeatures(nn.Module):
|
|
"""Learnable Fourier Feature encoding for positional encoding."""
|
|
|
|
def __init__(self, in_features: int, num_features: int, scale: float = 1.0):
|
|
super().__init__()
|
|
self.num_features = num_features
|
|
# Initialize frequency matrix B ~ N(0, scale^2)
|
|
self.B = nn.Parameter(torch.randn(in_features, num_features) * scale)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
x: Input coordinates [batch, in_features]
|
|
Returns:
|
|
Fourier features [batch, 2*num_features] (cos and sin)
|
|
"""
|
|
# x @ B: [batch, num_features]
|
|
proj = 2 * np.pi * (x @ self.B)
|
|
return torch.cat([torch.cos(proj), torch.sin(proj)], dim=-1)
|
|
|
|
|
|
class BioheatPINN(nn.Module):
|
|
"""
|
|
Physics-Informed Neural Network for Pennes Bioheat Equation.
|
|
|
|
Network architecture:
|
|
Input: (x, y, z, t) -> Fourier Features -> MLP -> Temperature T
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
fourier_features: int = 64,
|
|
fourier_scale: float = 2.0,
|
|
hidden_layers: list = None,
|
|
activation: str = 'tanh'
|
|
):
|
|
super().__init__()
|
|
|
|
if hidden_layers is None:
|
|
hidden_layers = [128, 128, 128, 64]
|
|
|
|
# Fourier feature encoding for 4D input (x, y, z, t)
|
|
self.fourier = FourierFeatures(4, fourier_features, fourier_scale)
|
|
|
|
# MLP layers
|
|
layers = []
|
|
in_dim = 2 * fourier_features # cos + sin features
|
|
|
|
for hidden_dim in hidden_layers:
|
|
layers.append(nn.Linear(in_dim, hidden_dim))
|
|
if activation == 'tanh':
|
|
layers.append(nn.Tanh())
|
|
elif activation == 'swish':
|
|
layers.append(nn.SiLU())
|
|
elif activation == 'gelu':
|
|
layers.append(nn.GELU())
|
|
in_dim = hidden_dim
|
|
|
|
# Output layer (temperature, no activation)
|
|
layers.append(nn.Linear(in_dim, 1))
|
|
|
|
self.mlp = nn.Sequential(*layers)
|
|
|
|
# Initialize weights using Xavier
|
|
self._init_weights()
|
|
|
|
def _init_weights(self):
|
|
for m in self.mlp.modules():
|
|
if isinstance(m, nn.Linear):
|
|
nn.init.xavier_uniform_(m.weight)
|
|
nn.init.zeros_(m.bias)
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor,
|
|
z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Forward pass: predict temperature at given coordinates.
|
|
|
|
Args:
|
|
x, y, z: Spatial coordinates [batch]
|
|
t: Time coordinate [batch]
|
|
|
|
Returns:
|
|
Temperature prediction [batch, 1]
|
|
"""
|
|
coords = torch.stack([x, y, z, t], dim=-1)
|
|
features = self.fourier(coords)
|
|
return self.mlp(features)
|
|
|
|
def predict(self, coords: torch.Tensor) -> torch.Tensor:
|
|
"""Predict temperature from stacked coordinates [batch, 4]."""
|
|
features = self.fourier(coords)
|
|
return self.mlp(features)
|
|
|
|
|
|
# =============================================================================
|
|
# Physics Loss Computation
|
|
# =============================================================================
|
|
|
|
class PennesBioheatLoss:
|
|
"""
|
|
Computes physics-informed loss for Pennes Bioheat Equation.
|
|
|
|
Loss = L_physics + lambda_bc * L_boundary + lambda_ic * L_initial
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
tissue: TissueProperties,
|
|
blood: BloodProperties = None,
|
|
probe_power: float = 20.0,
|
|
probe_position: Tuple[float, float, float] = (0.0, 0.0, 0.0),
|
|
probe_radius: float = 0.005,
|
|
domain_bounds: Tuple[float, float, float] = (0.1, 0.1, 0.1),
|
|
lambda_bc: float = 1.0,
|
|
lambda_ic: float = 1.0,
|
|
):
|
|
self.tissue = tissue
|
|
self.blood = blood or BloodProperties()
|
|
self.probe_power = probe_power
|
|
self.probe_position = torch.tensor(probe_position)
|
|
self.probe_radius = probe_radius
|
|
self.domain_bounds = domain_bounds
|
|
self.lambda_bc = lambda_bc
|
|
self.lambda_ic = lambda_ic
|
|
|
|
# Precompute constants
|
|
self.rho_c = tissue.density * tissue.specific_heat
|
|
self.k = tissue.conductivity
|
|
self.perfusion_coeff = (
|
|
tissue.perfusion * self.blood.density * self.blood.specific_heat
|
|
)
|
|
self.T_arterial = self.blood.arterial_temp
|
|
self.Q_m = tissue.metabolic_heat
|
|
|
|
def heat_source(self, x: torch.Tensor, y: torch.Tensor,
|
|
z: torch.Tensor) -> torch.Tensor:
|
|
"""Gaussian heat source from ablation probe."""
|
|
dx = x - self.probe_position[0]
|
|
dy = y - self.probe_position[1]
|
|
dz = z - self.probe_position[2]
|
|
r_sq = dx**2 + dy**2 + dz**2
|
|
|
|
# Gaussian distribution
|
|
sigma = self.probe_radius
|
|
amplitude = self.probe_power / ((2 * np.pi * sigma**2) ** 1.5)
|
|
return amplitude * torch.exp(-r_sq / (2 * sigma**2))
|
|
|
|
def physics_residual(
|
|
self,
|
|
model: BioheatPINN,
|
|
x: torch.Tensor,
|
|
y: torch.Tensor,
|
|
z: torch.Tensor,
|
|
t: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute PDE residual: rho*c*dT/dt - k*laplacian(T) - perfusion - Q_m - Q_s
|
|
|
|
Uses automatic differentiation to compute derivatives.
|
|
"""
|
|
# Enable gradients for coordinates
|
|
x = x.requires_grad_(True)
|
|
y = y.requires_grad_(True)
|
|
z = z.requires_grad_(True)
|
|
t = t.requires_grad_(True)
|
|
|
|
# Forward pass
|
|
T = model(x, y, z, t)
|
|
|
|
# First derivatives
|
|
grad_outputs = torch.ones_like(T)
|
|
|
|
dT_dx = torch.autograd.grad(T, x, grad_outputs, create_graph=True)[0]
|
|
dT_dy = torch.autograd.grad(T, y, grad_outputs, create_graph=True)[0]
|
|
dT_dz = torch.autograd.grad(T, z, grad_outputs, create_graph=True)[0]
|
|
dT_dt = torch.autograd.grad(T, t, grad_outputs, create_graph=True)[0]
|
|
|
|
# Second derivatives (Laplacian)
|
|
d2T_dx2 = torch.autograd.grad(dT_dx, x, grad_outputs, create_graph=True)[0]
|
|
d2T_dy2 = torch.autograd.grad(dT_dy, y, grad_outputs, create_graph=True)[0]
|
|
d2T_dz2 = torch.autograd.grad(dT_dz, z, grad_outputs, create_graph=True)[0]
|
|
|
|
laplacian_T = d2T_dx2 + d2T_dy2 + d2T_dz2
|
|
|
|
# Heat source
|
|
Q_s = self.heat_source(x, y, z)
|
|
|
|
# PDE residual
|
|
# rho*c*dT/dt = k*laplacian(T) + perfusion*(T_a - T) + Q_m + Q_s
|
|
time_term = self.rho_c * dT_dt
|
|
diffusion_term = self.k * laplacian_T
|
|
perfusion_term = self.perfusion_coeff * (self.T_arterial - T.squeeze())
|
|
source_term = self.Q_m + Q_s
|
|
|
|
residual = time_term - diffusion_term - perfusion_term - source_term
|
|
return residual
|
|
|
|
def boundary_loss(
|
|
self,
|
|
model: BioheatPINN,
|
|
x: torch.Tensor,
|
|
y: torch.Tensor,
|
|
z: torch.Tensor,
|
|
t: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Dirichlet boundary condition: T = T_body at boundaries."""
|
|
T = model(x, y, z, t)
|
|
T_body = self.T_arterial # Body temperature
|
|
return torch.mean((T.squeeze() - T_body) ** 2)
|
|
|
|
def initial_loss(
|
|
self,
|
|
model: BioheatPINN,
|
|
x: torch.Tensor,
|
|
y: torch.Tensor,
|
|
z: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Initial condition: T(x,y,z,0) = T_body everywhere."""
|
|
t = torch.zeros_like(x)
|
|
T = model(x, y, z, t)
|
|
T_body = self.T_arterial
|
|
return torch.mean((T.squeeze() - T_body) ** 2)
|
|
|
|
def compute_loss(
|
|
self,
|
|
model: BioheatPINN,
|
|
collocation_points: Dict[str, torch.Tensor],
|
|
boundary_points: Dict[str, torch.Tensor],
|
|
initial_points: Dict[str, torch.Tensor],
|
|
) -> Tuple[torch.Tensor, Dict[str, float]]:
|
|
"""
|
|
Compute total PINN loss.
|
|
|
|
Returns:
|
|
total_loss: Combined loss tensor
|
|
loss_dict: Dictionary with individual loss components
|
|
"""
|
|
# Physics loss at collocation points
|
|
residual = self.physics_residual(
|
|
model,
|
|
collocation_points['x'],
|
|
collocation_points['y'],
|
|
collocation_points['z'],
|
|
collocation_points['t'],
|
|
)
|
|
physics_loss = torch.mean(residual ** 2)
|
|
|
|
# Boundary loss
|
|
bc_loss = self.boundary_loss(
|
|
model,
|
|
boundary_points['x'],
|
|
boundary_points['y'],
|
|
boundary_points['z'],
|
|
boundary_points['t'],
|
|
)
|
|
|
|
# Initial condition loss
|
|
ic_loss = self.initial_loss(
|
|
model,
|
|
initial_points['x'],
|
|
initial_points['y'],
|
|
initial_points['z'],
|
|
)
|
|
|
|
# Total loss
|
|
total_loss = physics_loss + self.lambda_bc * bc_loss + self.lambda_ic * ic_loss
|
|
|
|
loss_dict = {
|
|
'total': total_loss.item(),
|
|
'physics': physics_loss.item(),
|
|
'boundary': bc_loss.item(),
|
|
'initial': ic_loss.item(),
|
|
}
|
|
|
|
return total_loss, loss_dict
|
|
|
|
|
|
# =============================================================================
|
|
# Training Loop
|
|
# =============================================================================
|
|
|
|
class BioheatTrainer:
|
|
"""Trainer for Pennes Bioheat PINN."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: BioheatPINN,
|
|
loss_fn: PennesBioheatLoss,
|
|
device: torch.device,
|
|
learning_rate: float = 1e-3,
|
|
num_collocation: int = 4096,
|
|
num_boundary: int = 1024,
|
|
num_initial: int = 1024,
|
|
domain_bounds: Tuple[float, float, float] = (0.1, 0.1, 0.1),
|
|
t_end: float = 600.0,
|
|
):
|
|
self.model = model.to(device)
|
|
self.loss_fn = loss_fn
|
|
self.device = device
|
|
self.optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
|
|
|
self.num_collocation = num_collocation
|
|
self.num_boundary = num_boundary
|
|
self.num_initial = num_initial
|
|
self.domain_bounds = domain_bounds
|
|
self.t_end = t_end
|
|
|
|
self.loss_history = []
|
|
self.step_count = 0
|
|
|
|
def sample_collocation_points(self) -> Dict[str, torch.Tensor]:
|
|
"""Sample random collocation points in domain."""
|
|
lx, ly, lz = self.domain_bounds
|
|
|
|
x = (torch.rand(self.num_collocation) - 0.5) * lx
|
|
y = (torch.rand(self.num_collocation) - 0.5) * ly
|
|
z = (torch.rand(self.num_collocation) - 0.5) * lz
|
|
t = torch.rand(self.num_collocation) * self.t_end
|
|
|
|
return {
|
|
'x': x.to(self.device),
|
|
'y': y.to(self.device),
|
|
'z': z.to(self.device),
|
|
't': t.to(self.device),
|
|
}
|
|
|
|
def sample_boundary_points(self) -> Dict[str, torch.Tensor]:
|
|
"""Sample points on domain boundary."""
|
|
lx, ly, lz = self.domain_bounds
|
|
n = self.num_boundary // 6 # 6 faces
|
|
|
|
points_x, points_y, points_z = [], [], []
|
|
|
|
# Sample from each face
|
|
for face_val, coord_idx in [
|
|
(-lx/2, 0), (lx/2, 0), # x faces
|
|
(-ly/2, 1), (ly/2, 1), # y faces
|
|
(-lz/2, 2), (lz/2, 2), # z faces
|
|
]:
|
|
if coord_idx == 0:
|
|
x = torch.full((n,), face_val)
|
|
y = (torch.rand(n) - 0.5) * ly
|
|
z = (torch.rand(n) - 0.5) * lz
|
|
elif coord_idx == 1:
|
|
x = (torch.rand(n) - 0.5) * lx
|
|
y = torch.full((n,), face_val)
|
|
z = (torch.rand(n) - 0.5) * lz
|
|
else:
|
|
x = (torch.rand(n) - 0.5) * lx
|
|
y = (torch.rand(n) - 0.5) * ly
|
|
z = torch.full((n,), face_val)
|
|
|
|
points_x.append(x)
|
|
points_y.append(y)
|
|
points_z.append(z)
|
|
|
|
x = torch.cat(points_x)
|
|
y = torch.cat(points_y)
|
|
z = torch.cat(points_z)
|
|
t = torch.rand(len(x)) * self.t_end
|
|
|
|
return {
|
|
'x': x.to(self.device),
|
|
'y': y.to(self.device),
|
|
'z': z.to(self.device),
|
|
't': t.to(self.device),
|
|
}
|
|
|
|
def sample_initial_points(self) -> Dict[str, torch.Tensor]:
|
|
"""Sample points at t=0 for initial condition."""
|
|
lx, ly, lz = self.domain_bounds
|
|
|
|
x = (torch.rand(self.num_initial) - 0.5) * lx
|
|
y = (torch.rand(self.num_initial) - 0.5) * ly
|
|
z = (torch.rand(self.num_initial) - 0.5) * lz
|
|
|
|
return {
|
|
'x': x.to(self.device),
|
|
'y': y.to(self.device),
|
|
'z': z.to(self.device),
|
|
}
|
|
|
|
def train_step(self) -> Dict[str, float]:
|
|
"""Execute one training step."""
|
|
self.model.train()
|
|
self.optimizer.zero_grad()
|
|
|
|
# Sample points
|
|
collocation = self.sample_collocation_points()
|
|
boundary = self.sample_boundary_points()
|
|
initial = self.sample_initial_points()
|
|
|
|
# Compute loss
|
|
loss, loss_dict = self.loss_fn.compute_loss(
|
|
self.model, collocation, boundary, initial
|
|
)
|
|
|
|
# Backprop
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
|
|
self.step_count += 1
|
|
loss_dict['step'] = self.step_count
|
|
self.loss_history.append(loss_dict)
|
|
|
|
return loss_dict
|
|
|
|
def train(self, num_steps: int, verbose: bool = True) -> list:
|
|
"""Train for multiple steps."""
|
|
losses = []
|
|
|
|
for i in range(num_steps):
|
|
loss_dict = self.train_step()
|
|
losses.append(loss_dict)
|
|
|
|
if verbose and (i + 1) % 100 == 0:
|
|
print(f"Step {self.step_count}: "
|
|
f"total={loss_dict['total']:.4e}, "
|
|
f"physics={loss_dict['physics']:.4e}, "
|
|
f"boundary={loss_dict['boundary']:.4e}")
|
|
|
|
return losses
|
|
|
|
|
|
# =============================================================================
|
|
# Benchmarking
|
|
# =============================================================================
|
|
|
|
def benchmark_pytorch(
|
|
device: str = 'cuda',
|
|
num_steps: int = 1000,
|
|
fourier_features: int = 64,
|
|
hidden_layers: list = None,
|
|
num_collocation: int = 4096,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Run benchmark for PyTorch Bioheat PINN.
|
|
|
|
Returns:
|
|
Dictionary with timing and performance metrics
|
|
"""
|
|
if hidden_layers is None:
|
|
hidden_layers = [128, 128, 128, 64]
|
|
|
|
# Setup device
|
|
if device == 'cuda' and not torch.cuda.is_available():
|
|
print("CUDA not available, falling back to CPU")
|
|
device = 'cpu'
|
|
|
|
device = torch.device(device)
|
|
print(f"Running benchmark on: {device}")
|
|
|
|
# Create model
|
|
model = BioheatPINN(
|
|
fourier_features=fourier_features,
|
|
hidden_layers=hidden_layers,
|
|
)
|
|
|
|
# Count parameters
|
|
num_params = sum(p.numel() for p in model.parameters())
|
|
print(f"Model parameters: {num_params:,}")
|
|
|
|
# Create loss function
|
|
tissue = TissueProperties.liver()
|
|
loss_fn = PennesBioheatLoss(
|
|
tissue=tissue,
|
|
probe_power=20.0,
|
|
probe_position=(0.0, 0.0, 0.0),
|
|
)
|
|
|
|
# Create trainer
|
|
trainer = BioheatTrainer(
|
|
model=model,
|
|
loss_fn=loss_fn,
|
|
device=device,
|
|
num_collocation=num_collocation,
|
|
)
|
|
|
|
# Warmup
|
|
print("Warming up...")
|
|
for _ in range(10):
|
|
trainer.train_step()
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
# Benchmark training
|
|
print(f"Benchmarking {num_steps} training steps...")
|
|
start_time = time.perf_counter()
|
|
|
|
trainer.train(num_steps, verbose=True)
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
training_time = time.perf_counter() - start_time
|
|
steps_per_second = num_steps / training_time
|
|
|
|
# Benchmark inference
|
|
print("Benchmarking inference...")
|
|
model.eval()
|
|
|
|
# Create inference grid (32^3 = 32768 points)
|
|
grid_size = 32
|
|
x = torch.linspace(-0.05, 0.05, grid_size)
|
|
y = torch.linspace(-0.05, 0.05, grid_size)
|
|
z = torch.linspace(-0.05, 0.05, grid_size)
|
|
xx, yy, zz = torch.meshgrid(x, y, z, indexing='ij')
|
|
|
|
coords = torch.stack([
|
|
xx.flatten(),
|
|
yy.flatten(),
|
|
zz.flatten(),
|
|
torch.full((grid_size**3,), 300.0), # t = 300s
|
|
], dim=-1).to(device)
|
|
|
|
# Warmup inference
|
|
with torch.no_grad():
|
|
for _ in range(5):
|
|
_ = model.predict(coords)
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
# Time inference
|
|
num_inference_runs = 100
|
|
start_time = time.perf_counter()
|
|
|
|
with torch.no_grad():
|
|
for _ in range(num_inference_runs):
|
|
_ = model.predict(coords)
|
|
|
|
if device.type == 'cuda':
|
|
torch.cuda.synchronize()
|
|
|
|
inference_time = time.perf_counter() - start_time
|
|
inference_ms = (inference_time / num_inference_runs) * 1000
|
|
|
|
# Memory usage
|
|
if device.type == 'cuda':
|
|
peak_memory_mb = torch.cuda.max_memory_allocated() / (1024 ** 2)
|
|
else:
|
|
peak_memory_mb = 0.0
|
|
|
|
# Final loss
|
|
final_loss = trainer.loss_history[-1]['total']
|
|
|
|
results = {
|
|
'device': str(device),
|
|
'num_parameters': num_params,
|
|
'num_steps': num_steps,
|
|
'training_time_sec': training_time,
|
|
'steps_per_second': steps_per_second,
|
|
'inference_grid_size': grid_size ** 3,
|
|
'inference_mean_ms': inference_ms,
|
|
'peak_memory_mb': peak_memory_mb,
|
|
'final_loss': final_loss,
|
|
}
|
|
|
|
return results
|
|
|
|
|
|
def print_benchmark_results(results: Dict[str, Any]):
|
|
"""Pretty print benchmark results."""
|
|
print("\n" + "=" * 60)
|
|
print("PyTorch Bioheat PINN Benchmark Results")
|
|
print("=" * 60)
|
|
print(f"Device: {results['device']}")
|
|
print(f"Parameters: {results['num_parameters']:,}")
|
|
print(f"Training steps: {results['num_steps']}")
|
|
print("-" * 60)
|
|
print(f"Training time: {results['training_time_sec']:.2f} sec")
|
|
print(f"Steps/second: {results['steps_per_second']:.1f}")
|
|
print(f"Inference (32^3): {results['inference_mean_ms']:.2f} ms")
|
|
if results['peak_memory_mb'] > 0:
|
|
print(f"Peak GPU memory: {results['peak_memory_mb']:.1f} MB")
|
|
print(f"Final loss: {results['final_loss']:.4e}")
|
|
print("=" * 60)
|
|
|
|
|
|
# =============================================================================
|
|
# Main
|
|
# =============================================================================
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description='PyTorch Reference Implementation for Pennes Bioheat PINN'
|
|
)
|
|
parser.add_argument(
|
|
'--device', type=str, default='cuda',
|
|
choices=['cuda', 'cpu'],
|
|
help='Device to run on (default: cuda)'
|
|
)
|
|
parser.add_argument(
|
|
'--steps', type=int, default=1000,
|
|
help='Number of training steps (default: 1000)'
|
|
)
|
|
parser.add_argument(
|
|
'--fourier-features', type=int, default=64,
|
|
help='Number of Fourier features (default: 64)'
|
|
)
|
|
parser.add_argument(
|
|
'--collocation', type=int, default=4096,
|
|
help='Number of collocation points (default: 4096)'
|
|
)
|
|
parser.add_argument(
|
|
'--output', type=str, default=None,
|
|
help='Output JSON file for results'
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
results = benchmark_pytorch(
|
|
device=args.device,
|
|
num_steps=args.steps,
|
|
fourier_features=args.fourier_features,
|
|
num_collocation=args.collocation,
|
|
)
|
|
|
|
print_benchmark_results(results)
|
|
|
|
if args.output:
|
|
import json
|
|
with open(args.output, 'w') as f:
|
|
json.dump(results, f, indent=2)
|
|
print(f"\nResults saved to: {args.output}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|