72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Debug script to compare lifting MLP alone.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
def debug_lifting_mlp():
|
|
"""Create a minimal test case for lifting MLP comparison."""
|
|
torch.manual_seed(42)
|
|
|
|
# Configuration matching the FNO
|
|
in_channels = 3 # After positional encoding: data_channels(1) + 2
|
|
width = 32
|
|
hidden_dim = 2 * width # 64
|
|
height, width_spatial = 8, 8
|
|
|
|
# Create input: [batch, in_channels, height, width]
|
|
x = torch.randn(1, in_channels, height, width_spatial)
|
|
|
|
# Create weights
|
|
scale_fc1 = np.sqrt(2.0 / (in_channels + hidden_dim))
|
|
scale_fc2 = np.sqrt(2.0 / (hidden_dim + width))
|
|
|
|
fc1_weight = torch.randn(hidden_dim, in_channels) * scale_fc1
|
|
fc1_bias = torch.zeros(hidden_dim)
|
|
fc2_weight = torch.randn(width, hidden_dim) * scale_fc2
|
|
fc2_bias = torch.zeros(width)
|
|
|
|
print(f"Input shape: {x.shape}")
|
|
print(f"FC1 weight shape: {fc1_weight.shape}")
|
|
print(f"FC2 weight shape: {fc2_weight.shape}")
|
|
|
|
# Apply lifting MLP (same as in run_fno_forward_manual)
|
|
# Reshape for linear: [B, C, H, W] -> [B, H, W, C]
|
|
x_perm = x.permute(0, 2, 3, 1)
|
|
print(f"After permute: {x_perm.shape}")
|
|
|
|
# FC1 with GELU
|
|
x1 = torch.nn.functional.linear(x_perm, fc1_weight, fc1_bias)
|
|
print(f"After FC1: {x1.shape}, range=[{x1.min():.4f}, {x1.max():.4f}]")
|
|
|
|
x1_gelu = torch.nn.functional.gelu(x1)
|
|
print(f"After GELU: {x1_gelu.shape}, range=[{x1_gelu.min():.4f}, {x1_gelu.max():.4f}]")
|
|
|
|
# FC2
|
|
x2 = torch.nn.functional.linear(x1_gelu, fc2_weight, fc2_bias)
|
|
print(f"After FC2: {x2.shape}, range=[{x2.min():.4f}, {x2.max():.4f}]")
|
|
|
|
# Reshape back: [B, H, W, C] -> [B, C, H, W]
|
|
output = x2.permute(0, 3, 1, 2)
|
|
print(f"Final output: {output.shape}, range=[{output.min():.4f}, {output.max():.4f}]")
|
|
|
|
# Save for Rust comparison
|
|
save_file({
|
|
"input": x.contiguous(),
|
|
"fc1_weight": fc1_weight.contiguous(),
|
|
"fc1_bias": fc1_bias.contiguous(),
|
|
"fc2_weight": fc2_weight.contiguous(),
|
|
"fc2_bias": fc2_bias.contiguous(),
|
|
"output": output.contiguous(),
|
|
}, "/tmp/lifting_mlp_debug.safetensors")
|
|
|
|
print("\nSaved debug data to /tmp/lifting_mlp_debug.safetensors")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
debug_lifting_mlp()
|