From fee3ba7be896ed767d1ae680a0811e023f71fce1 Mon Sep 17 00:00:00 2001 From: skal Date: Fri, 13 Feb 2026 16:42:36 +0100 Subject: CNN v2 export: Read and display mip_level from checkpoints Export scripts now read mip_level from checkpoint config and display it. Shader generator includes mip level in generated comments. Changes: - export_cnn_v2_weights.py: Read mip_level, print in config - export_cnn_v2_shader.py: Read mip_level, pass to shader gen, add to comments Co-Authored-By: Claude Sonnet 4.5 --- training/export_cnn_v2_shader.py | 7 ++++++- training/export_cnn_v2_weights.py | 2 ++ 2 files changed, 8 insertions(+), 1 deletion(-) (limited to 'training') diff --git a/training/export_cnn_v2_shader.py b/training/export_cnn_v2_shader.py index dc475d8..1c74ad0 100755 --- a/training/export_cnn_v2_shader.py +++ b/training/export_cnn_v2_shader.py @@ -14,7 +14,7 @@ import torch from pathlib import Path -def export_layer_shader(layer_idx, weights, kernel_size, output_dir, is_output_layer=False): +def export_layer_shader(layer_idx, weights, kernel_size, output_dir, mip_level=0, is_output_layer=False): """Generate WGSL compute shader for a single CNN layer. Args: @@ -22,6 +22,7 @@ def export_layer_shader(layer_idx, weights, kernel_size, output_dir, is_output_l weights: (4, 12, k, k) weight tensor (uniform 12D→4D) kernel_size: Kernel size (3, 5, etc.) output_dir: Output directory path + mip_level: Mip level used for p0-p3 (0=original, 1=half, etc.) is_output_layer: True if this is the final RGBA output layer """ weights_flat = weights.flatten() @@ -44,6 +45,7 @@ def export_layer_shader(layer_idx, weights, kernel_size, output_dir, is_output_l shader_code = f"""// CNN v2 Layer {layer_idx} - Auto-generated (uniform 12D→4D) // Kernel: {kernel_size}×{kernel_size}, In: 12D (4 prev + 8 static), Out: 4D +// Mip level: {mip_level} (p0-p3 features) const KERNEL_SIZE: u32 = {kernel_size}u; const IN_CHANNELS: u32 = 12u; // 4 (input/prev) + 8 (static) @@ -164,10 +166,12 @@ def export_checkpoint(checkpoint_path, output_dir): kernel_size = config.get('kernel_size', 3) num_layers = config.get('num_layers', 3) + mip_level = config.get('mip_level', 0) print(f"Configuration:") print(f" Kernel size: {kernel_size}×{kernel_size}") print(f" Layers: {num_layers}") + print(f" Mip level: {mip_level} (p0-p3 features)") print(f" Architecture: uniform 12D→4D") output_dir = Path(output_dir) @@ -189,6 +193,7 @@ def export_checkpoint(checkpoint_path, output_dir): weights=layer_weights, kernel_size=kernel_size, output_dir=output_dir, + mip_level=mip_level, is_output_layer=is_output ) diff --git a/training/export_cnn_v2_weights.py b/training/export_cnn_v2_weights.py index bbe94dd..9e9e352 100755 --- a/training/export_cnn_v2_weights.py +++ b/training/export_cnn_v2_weights.py @@ -56,10 +56,12 @@ def export_weights_binary(checkpoint_path, output_path): kernel_sizes = [3, 3, 3] # fallback num_layers = config.get('num_layers', len(kernel_sizes)) + mip_level = config.get('mip_level', 0) print(f"Configuration:") print(f" Kernel sizes: {kernel_sizes}") print(f" Layers: {num_layers}") + print(f" Mip level: {mip_level} (p0-p3 features)") print(f" Architecture: uniform 12D→4D (bias=False)") # Collect layer info - all layers uniform 12D→4D -- cgit v1.2.3