200 lines
7.2 KiB
Python
200 lines
7.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Debug script to trace through FNO forward pass step by step.
|
|
Saves intermediate outputs at each stage for comparison with Rust.
|
|
"""
|
|
|
|
import json
|
|
import numpy as np
|
|
import torch
|
|
from pathlib import Path
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
def create_grid_positional_encoding(height: int, width: int) -> torch.Tensor:
|
|
"""Create 2D grid positional encoding."""
|
|
y = torch.linspace(-1, 1, height)
|
|
x = torch.linspace(-1, 1, width)
|
|
grid_y, grid_x = torch.meshgrid(y, x, indexing='ij')
|
|
grid = torch.stack([grid_x, grid_y], dim=0)
|
|
return grid
|
|
|
|
|
|
def debug_fno_step_by_step():
|
|
"""Run FNO forward pass with intermediate outputs."""
|
|
torch.manual_seed(42)
|
|
|
|
# Simple config
|
|
batch_size = 1
|
|
data_channels = 1
|
|
height = 32
|
|
width = 32
|
|
model_width = 16 # Smaller for faster debugging
|
|
n_modes = (8, 8)
|
|
n_layers = 2 # Fewer layers
|
|
|
|
# Create deterministic input
|
|
y = torch.linspace(0, 2 * np.pi, height)
|
|
x = torch.linspace(0, 2 * np.pi, width)
|
|
yy, xx = torch.meshgrid(y, x, indexing='ij')
|
|
input_tensor = torch.sin(xx) * torch.cos(yy)
|
|
input_tensor = input_tensor.unsqueeze(0).unsqueeze(0) # [1, 1, H, W]
|
|
|
|
print(f"Input: shape={input_tensor.shape}, range=[{input_tensor.min():.4f}, {input_tensor.max():.4f}]")
|
|
|
|
# Step 1: Positional encoding
|
|
grid = create_grid_positional_encoding(height, width)
|
|
grid = grid.unsqueeze(0).expand(batch_size, -1, -1, -1)
|
|
with_pos = torch.cat([input_tensor, grid], dim=1)
|
|
print(f"After pos encoding: shape={with_pos.shape}, range=[{with_pos.min():.4f}, {with_pos.max():.4f}]")
|
|
|
|
# Create weights
|
|
in_with_pos = data_channels + 2
|
|
hidden_dim = 2 * model_width
|
|
|
|
# Lifting weights
|
|
scale_fc1 = np.sqrt(2.0 / (in_with_pos + hidden_dim))
|
|
scale_fc2 = np.sqrt(2.0 / (hidden_dim + model_width))
|
|
lifting_fc1_weight = torch.randn(hidden_dim, in_with_pos) * scale_fc1
|
|
lifting_fc1_bias = torch.zeros(hidden_dim)
|
|
lifting_fc2_weight = torch.randn(model_width, hidden_dim) * scale_fc2
|
|
lifting_fc2_bias = torch.zeros(model_width)
|
|
|
|
# Step 2: Lifting MLP
|
|
x = with_pos.permute(0, 2, 3, 1) # [B, H, W, C]
|
|
x = torch.nn.functional.linear(x, lifting_fc1_weight, lifting_fc1_bias)
|
|
x = torch.nn.functional.gelu(x)
|
|
x = torch.nn.functional.linear(x, lifting_fc2_weight, lifting_fc2_bias)
|
|
x = x.permute(0, 3, 1, 2) # [B, C, H, W]
|
|
print(f"After lifting: shape={x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")
|
|
after_lifting = x.clone()
|
|
|
|
# Spectral and skip weights
|
|
spectral_weights = []
|
|
skip_weights = []
|
|
scale_spectral = 1.0 / (model_width * n_modes[0] * n_modes[1])
|
|
scale_conv = np.sqrt(2.0 / (model_width + model_width))
|
|
|
|
for i in range(n_layers):
|
|
sw = {
|
|
'w1_real': torch.randn(model_width, model_width, n_modes[0], n_modes[1]) * scale_spectral,
|
|
'w1_imag': torch.randn(model_width, model_width, n_modes[0], n_modes[1]) * scale_spectral,
|
|
'w2_real': torch.randn(model_width, model_width, n_modes[0], n_modes[1]) * scale_spectral,
|
|
'w2_imag': torch.randn(model_width, model_width, n_modes[0], n_modes[1]) * scale_spectral,
|
|
}
|
|
spectral_weights.append(sw)
|
|
skip_weights.append({
|
|
'weight': torch.randn(model_width, model_width) * scale_conv,
|
|
'bias': torch.zeros(model_width),
|
|
})
|
|
|
|
# Step 3: Fourier layers
|
|
for i in range(n_layers):
|
|
# Spectral conv path
|
|
x_ft = torch.fft.fft2(x)
|
|
|
|
sw = spectral_weights[i]
|
|
w1 = torch.complex(sw['w1_real'], sw['w1_imag'])
|
|
w2 = torch.complex(sw['w2_real'], sw['w2_imag'])
|
|
|
|
modes_h, modes_w = n_modes
|
|
out_ft = torch.zeros_like(x_ft)
|
|
|
|
x_upper = x_ft[:, :, :modes_h, :modes_w]
|
|
out_upper = torch.einsum('bihw,iohw->bohw', x_upper, w1)
|
|
out_ft[:, :, :modes_h, :modes_w] = out_upper
|
|
|
|
x_lower = x_ft[:, :, -modes_h:, :modes_w]
|
|
out_lower = torch.einsum('bihw,iohw->bohw', x_lower, w2)
|
|
out_ft[:, :, -modes_h:, :modes_w] = out_lower
|
|
|
|
x_spectral = torch.fft.ifft2(out_ft).real
|
|
print(f" Layer {i} spectral: range=[{x_spectral.min():.4f}, {x_spectral.max():.4f}]")
|
|
|
|
# Skip connection path
|
|
skip_w = skip_weights[i]['weight']
|
|
skip_b = skip_weights[i]['bias']
|
|
x_skip = x.permute(0, 2, 3, 1)
|
|
x_skip = torch.nn.functional.linear(x_skip, skip_w, skip_b)
|
|
x_skip = x_skip.permute(0, 3, 1, 2)
|
|
print(f" Layer {i} skip: range=[{x_skip.min():.4f}, {x_skip.max():.4f}]")
|
|
|
|
# Combine
|
|
x = x_spectral + x_skip
|
|
print(f" Layer {i} combined: range=[{x.min():.4f}, {x.max():.4f}]")
|
|
|
|
if i < n_layers - 1:
|
|
x = torch.nn.functional.gelu(x)
|
|
print(f" Layer {i} after GELU: range=[{x.min():.4f}, {x.max():.4f}]")
|
|
|
|
after_fourier = x.clone()
|
|
print(f"After all Fourier layers: shape={x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")
|
|
|
|
# Step 4: Projection
|
|
proj_hidden = 128
|
|
scale_proj1 = np.sqrt(2.0 / (model_width + proj_hidden))
|
|
scale_proj2 = np.sqrt(2.0 / (proj_hidden + 1))
|
|
|
|
proj1_weight = torch.randn(proj_hidden, model_width) * scale_proj1
|
|
proj1_bias = torch.zeros(proj_hidden)
|
|
proj2_weight = torch.randn(1, proj_hidden) * scale_proj2
|
|
proj2_bias = torch.zeros(1)
|
|
|
|
x = x.permute(0, 2, 3, 1)
|
|
x = torch.nn.functional.linear(x, proj1_weight, proj1_bias)
|
|
x = torch.nn.functional.gelu(x)
|
|
x = torch.nn.functional.linear(x, proj2_weight, proj2_bias)
|
|
x = x.permute(0, 3, 1, 2)
|
|
|
|
print(f"Final output: shape={x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")
|
|
print(f"Final output mean: {x.mean():.6f}")
|
|
|
|
# Save all data
|
|
output_dir = Path("/tmp/fno_debug")
|
|
output_dir.mkdir(exist_ok=True)
|
|
|
|
tensors = {
|
|
"input": input_tensor.contiguous(),
|
|
"after_pos_encoding": with_pos.contiguous(),
|
|
"lifting_fc1_weight": lifting_fc1_weight.contiguous(),
|
|
"lifting_fc1_bias": lifting_fc1_bias.contiguous(),
|
|
"lifting_fc2_weight": lifting_fc2_weight.contiguous(),
|
|
"lifting_fc2_bias": lifting_fc2_bias.contiguous(),
|
|
"after_lifting": after_lifting.contiguous(),
|
|
"after_fourier": after_fourier.contiguous(),
|
|
"proj1_weight": proj1_weight.contiguous(),
|
|
"proj1_bias": proj1_bias.contiguous(),
|
|
"proj2_weight": proj2_weight.contiguous(),
|
|
"proj2_bias": proj2_bias.contiguous(),
|
|
"output": x.contiguous(),
|
|
}
|
|
|
|
# Add spectral and skip weights
|
|
for i in range(n_layers):
|
|
sw = spectral_weights[i]
|
|
for key, tensor in sw.items():
|
|
tensors[f"spectral_{i}_{key}"] = tensor.contiguous()
|
|
tensors[f"skip_{i}_weight"] = skip_weights[i]['weight'].contiguous()
|
|
tensors[f"skip_{i}_bias"] = skip_weights[i]['bias'].contiguous()
|
|
|
|
save_file(tensors, output_dir / "debug.safetensors")
|
|
|
|
# Save config
|
|
config = {
|
|
"batch_size": batch_size,
|
|
"data_channels": data_channels,
|
|
"height": height,
|
|
"width": width,
|
|
"model_width": model_width,
|
|
"n_modes": n_modes,
|
|
"n_layers": n_layers,
|
|
}
|
|
with open(output_dir / "config.json", "w") as f:
|
|
json.dump(config, f, indent=2)
|
|
|
|
print(f"\nSaved debug data to {output_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
debug_fno_step_by_step()
|