From b25132133590e39967ebd0f3205123ee6628674b Mon Sep 17 00:00:00 2001 From: skal Date: Tue, 10 Feb 2026 21:21:53 +0100 Subject: fix: Add clamp to CNN final layer to match PyTorch training MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CNN output mismatch resolved: final layer (7→1) now clamps to [0,1]. Changes: - Add clamp(sum, 0.0, 1.0) to cnn_conv3x3_7to1 and cnn_conv5x5_7to1 - Add generate_conv_final_function() to train_cnn.py for auto-generation - Update comments to clarify clamping behavior - Future exports will auto-generate final layers with correct clamp PyTorch uses torch.clamp(out, 0.0, 1.0) on final output; shaders were missing this critical operation, causing range mismatches. Co-Authored-By: Claude Sonnet 4.5 --- workspaces/main/shaders/cnn/cnn_conv3x3.wgsl | 2 +- workspaces/main/shaders/cnn/cnn_conv5x5.wgsl | 2 +- workspaces/main/shaders/cnn/cnn_layer.wgsl | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) (limited to 'workspaces/main') diff --git a/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl b/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl index 79b0350..00eae22 100644 --- a/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl +++ b/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl @@ -138,5 +138,5 @@ fn cnn_conv3x3_7to1( } } - return sum; // Output in [-1,1] + return clamp(sum, 0.0, 1.0); // Match PyTorch clamp } diff --git a/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl b/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl index 5570589..4f0a5f3 100644 --- a/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl +++ b/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl @@ -77,7 +77,7 @@ fn cnn_conv5x5_7to1( } } - return sum; // Output in [-1,1] + return clamp(sum, 0.0, 1.0); // Match PyTorch clamp } // Source layer: 7→4 channels (RGBD output) diff --git a/workspaces/main/shaders/cnn/cnn_layer.wgsl b/workspaces/main/shaders/cnn/cnn_layer.wgsl index e67ad31..8eccb26 100644 --- a/workspaces/main/shaders/cnn/cnn_layer.wgsl +++ b/workspaces/main/shaders/cnn/cnn_layer.wgsl @@ -49,7 +49,7 @@ struct CNNLayerParams { else if (params.layer_index == 2) { let gray_out = cnn_conv3x3_7to1(txt, smplr, uv, uniforms.resolution, gray, weights_layer2); - // gray_out already in [0,1] from clipped training + // gray_out in [0,1] (clamped to match PyTorch training) result = vec4(gray_out, gray_out, gray_out, 1.0); return mix(original_raw, result, params.blend_amount); // [0,1] } -- cgit v1.2.3