From 161a59fa50bb92e3664c389fa03b95aefe349b3f Mon Sep 17 00:00:00 2001 From: skal Date: Sun, 15 Feb 2026 18:44:17 +0100 Subject: refactor(cnn): isolate CNN v2 to cnn_v2/ subdirectory Move all CNN v2 files to dedicated cnn_v2/ directory to prepare for CNN v3 development. Zero functional changes. Structure: - cnn_v2/src/ - C++ effect implementation - cnn_v2/shaders/ - WGSL shaders (6 files) - cnn_v2/weights/ - Binary weights (3 files) - cnn_v2/training/ - Python training scripts (4 files) - cnn_v2/scripts/ - Shell scripts (train_cnn_v2_full.sh) - cnn_v2/tools/ - Validation tools (HTML) - cnn_v2/docs/ - Documentation (4 markdown files) Changes: - Update CMake source list to cnn_v2/src/cnn_v2_effect.cc - Update assets.txt with relative paths to cnn_v2/ - Update includes to ../../cnn_v2/src/cnn_v2_effect.h - Add PROJECT_ROOT resolution to Python/shell scripts - Update doc references in HOWTO.md, TODO.md - Add cnn_v2/README.md Verification: 34/34 tests passing, demo runs correctly. Co-Authored-By: Claude Sonnet 4.5 --- cnn_v2/README.md | 60 + cnn_v2/docs/CNN_V2.md | 813 ++++++++++++ cnn_v2/docs/CNN_V2_BINARY_FORMAT.md | 235 ++++ cnn_v2/docs/CNN_V2_DEBUG_TOOLS.md | 143 ++ cnn_v2/docs/CNN_V2_WEB_TOOL.md | 348 +++++ cnn_v2/scripts/train_cnn_v2_full.sh | 428 ++++++ cnn_v2/shaders/cnn_v2_compute.wgsl | 143 ++ cnn_v2/shaders/cnn_v2_layer_0.wgsl | 174 +++ cnn_v2/shaders/cnn_v2_layer_1.wgsl | 174 +++ cnn_v2/shaders/cnn_v2_layer_2.wgsl | 156 +++ cnn_v2/shaders/cnn_v2_layer_template.wgsl | 68 + cnn_v2/shaders/cnn_v2_static.wgsl | 75 ++ cnn_v2/src/cnn_v2_effect.cc | 497 +++++++ cnn_v2/src/cnn_v2_effect.h | 89 ++ cnn_v2/tools/cnn_v2_test/README.md | 251 ++++ cnn_v2/tools/cnn_v2_test/index.html | 2014 +++++++++++++++++++++++++++++ cnn_v2/training/export_cnn_v2_shader.py | 218 ++++ cnn_v2/training/export_cnn_v2_weights.py | 288 +++++ cnn_v2/training/gen_identity_weights.py | 175 +++ cnn_v2/training/train_cnn_v2.py | 472 +++++++ 20 files changed, 6821 insertions(+) create mode 100644 cnn_v2/README.md create mode 100644 cnn_v2/docs/CNN_V2.md create mode 100644 cnn_v2/docs/CNN_V2_BINARY_FORMAT.md create mode 100644 cnn_v2/docs/CNN_V2_DEBUG_TOOLS.md create mode 100644 cnn_v2/docs/CNN_V2_WEB_TOOL.md create mode 100755 cnn_v2/scripts/train_cnn_v2_full.sh create mode 100644 cnn_v2/shaders/cnn_v2_compute.wgsl create mode 100644 cnn_v2/shaders/cnn_v2_layer_0.wgsl create mode 100644 cnn_v2/shaders/cnn_v2_layer_1.wgsl create mode 100644 cnn_v2/shaders/cnn_v2_layer_2.wgsl create mode 100644 cnn_v2/shaders/cnn_v2_layer_template.wgsl create mode 100644 cnn_v2/shaders/cnn_v2_static.wgsl create mode 100644 cnn_v2/src/cnn_v2_effect.cc create mode 100644 cnn_v2/src/cnn_v2_effect.h create mode 100644 cnn_v2/tools/cnn_v2_test/README.md create mode 100644 cnn_v2/tools/cnn_v2_test/index.html create mode 100755 cnn_v2/training/export_cnn_v2_shader.py create mode 100755 cnn_v2/training/export_cnn_v2_weights.py create mode 100755 cnn_v2/training/gen_identity_weights.py create mode 100755 cnn_v2/training/train_cnn_v2.py (limited to 'cnn_v2') diff --git a/cnn_v2/README.md b/cnn_v2/README.md new file mode 100644 index 0000000..ef0cf44 --- /dev/null +++ b/cnn_v2/README.md @@ -0,0 +1,60 @@ +# CNN v2: Parametric Post-Processing Neural Network + +**Architecture:** 3-layer compute, storage buffer (~3.2 KB) +**Features:** 7D static (RGBD + UV + sin + bias), sigmoid activation + +## Quick Start + +```bash +./cnn_v2/scripts/train_cnn_v2_full.sh +``` + +## Documentation + +- [CNN_V2.md](docs/CNN_V2.md) - Architecture and implementation details +- [CNN_V2_BINARY_FORMAT.md](docs/CNN_V2_BINARY_FORMAT.md) - Weight format specification +- [CNN_V2_WEB_TOOL.md](docs/CNN_V2_WEB_TOOL.md) - Validation tool documentation +- [CNN_V2_DEBUG_TOOLS.md](docs/CNN_V2_DEBUG_TOOLS.md) - Debugging and analysis tools + +## Integration + +- **C++:** `cnn_v2/src/cnn_v2_effect.{h,cc}` +- **Assets:** `workspaces/main/assets.txt` (lines 47-49) +- **Test:** `src/tests/gpu/test_demo_effects.cc` (line 93) + +## Directory Structure + +``` +cnn_v2/ +├── README.md # This file +├── src/ +│ ├── cnn_v2_effect.h # Effect header +│ └── cnn_v2_effect.cc # Effect implementation +├── shaders/ # WGSL shaders (6 files) +├── weights/ # Binary weights (3 files) +├── training/ # Python training scripts (4 files) +├── scripts/ # Shell scripts (train_cnn_v2_full.sh) +├── tools/ # Validation tools (HTML) +└── docs/ # Documentation (4 markdown files) +``` + +## Training Pipeline + +1. **Train model:** `./cnn_v2/scripts/train_cnn_v2_full.sh` +2. **Export weights:** Automatic (binary format, ~3.2 KB) +3. **Validate:** HTML tool at `cnn_v2/tools/cnn_v2_test/index.html` + +For detailed training options: `./cnn_v2/scripts/train_cnn_v2_full.sh --help` + +## Key Features + +- **Parametric static features:** 7D input (RGBD + UV + sin encoding + bias) +- **Storage buffer architecture:** Dynamic layer count, compact binary format +- **Sigmoid activation:** Smooth gradients, prevents training collapse +- **Patch-based training:** Sample-efficient, focuses on salient regions +- **Sub-10KB target:** Achieved with 3-layer model (~3.2 KB) + +## Next Steps + +- **8-bit quantization:** 2× size reduction (~1.6 KB) via quantization-aware training (QAT) +- **CNN v3:** U-Net architecture for enhanced quality (separate directory) diff --git a/cnn_v2/docs/CNN_V2.md b/cnn_v2/docs/CNN_V2.md new file mode 100644 index 0000000..b7fd6f8 --- /dev/null +++ b/cnn_v2/docs/CNN_V2.md @@ -0,0 +1,813 @@ +# CNN v2: Parametric Static Features + +**Technical Design Document** + +--- + +## Overview + +CNN v2 extends the original CNN post-processing effect with parametric static features, enabling richer spatial and frequency-domain inputs for improved visual quality. + +**Key improvements over v1:** +- 7D static feature input (vs 4D RGB) +- Multi-frequency position encoding (NeRF-style) +- Configurable mip-level for p0-p3 parametric features (0-3) +- Per-layer configurable kernel sizes (1×1, 3×3, 5×5) +- Variable channel counts per layer +- Float16 weight storage (~3.2 KB for 3-layer model) +- Bias integrated as static feature dimension +- Storage buffer architecture (dynamic layer count) +- Binary weight format v2 for runtime loading +- Sigmoid activation for layer 0 and final layer (smooth [0,1] mapping) + +**Status:** ✅ Complete. Sigmoid activation, stable training, validation tools operational. + +**Breaking Change:** +- Models trained with `clamp()` incompatible. Retrain required. + +**TODO:** +- 8-bit quantization with QAT for 2× size reduction (~1.6 KB) + +--- + +## Architecture + +### Pipeline Overview + +``` +Input RGBD → Static Features Compute → CNN Layers → Output RGBA + └─ computed once/frame ─┘ └─ multi-pass ─┘ +``` + +**Detailed Data Flow:** + +``` + ┌─────────────────────────────────────────┐ + │ Static Features (computed once) │ + │ 8D: p0,p1,p2,p3,uv_x,uv_y,sin10x,bias │ + └──────────────┬──────────────────────────┘ + │ + │ 8D (broadcast to all layers) + ├───────────────────────────┐ + │ │ + ┌──────────────┐ │ │ + │ Input RGBD │──────────────┤ │ + │ 4D │ 4D │ │ + └──────────────┘ │ │ + ▼ │ + ┌────────────┐ │ + │ Layer 0 │ (12D input) │ + │ (CNN) │ = 4D + 8D │ + │ 12D → 4D │ │ + └─────┬──────┘ │ + │ 4D output │ + │ │ + ├───────────────────────────┘ + │ │ + ▼ │ + ┌────────────┐ │ + │ Layer 1 │ (12D input) │ + │ (CNN) │ = 4D + 8D │ + │ 12D → 4D │ │ + └─────┬──────┘ │ + │ 4D output │ + │ │ + ├───────────────────────────┘ + ▼ │ + ... │ + │ │ + ▼ │ + ┌────────────┐ │ + │ Layer N │ (12D input) │ + │ (output) │◄──────────────────┘ + │ 12D → 4D │ + └─────┬──────┘ + │ 4D (RGBA) + ▼ + Output +``` + +**Key Points:** +- Static features computed once, broadcast to all CNN layers +- Each layer: previous 4D output + 8D static → 12D input → 4D output +- Ping-pong buffering between layers +- Layer 0 special case: uses input RGBD instead of previous layer output + +**Static Features Texture:** +- Name: `static_features` +- Format: `texture_storage_2d` (4×u32) +- Data: 8 float16 values packed via `pack2x16float()` +- Computed once per frame, read by all CNN layers +- Lifetime: Entire frame (all CNN layer passes) + +**CNN Layers:** +- Layer 0: input RGBD (4D) + static (8D) = 12D → 4 channels +- Layer 1+: previous output (4D) + static (8D) = 12D → 4 channels +- All layers: uniform 12D input, 4D output (ping-pong buffer) +- Storage: `texture_storage_2d` (4 channels as 2×f16 pairs) + +**Activation Functions:** +- Layer 0 & final layer: `sigmoid(x)` for smooth [0,1] mapping +- Middle layers: `ReLU` (max(0, x)) +- Rationale: Sigmoid prevents gradient blocking at boundaries, enabling better convergence +- Breaking change: Models trained with `clamp(x, 0, 1)` are incompatible, retrain required + +--- + +## Static Features (7D + 1 bias) + +### Feature Layout + +**8 float16 values per pixel:** + +```wgsl +// Slot 0-3: Parametric features (p0, p1, p2, p3) +// Sampled from configurable mip level (0=original, 1=half, 2=quarter, 3=eighth) +// Training sets mip_level via --mip-level flag, stored in binary format v2 +let p0 = ...; // RGB.r from selected mip level +let p1 = ...; // RGB.g from selected mip level +let p2 = ...; // RGB.b from selected mip level +let p3 = ...; // Depth or RGB channel from mip level + +// Slot 4-5: UV coordinates (normalized screen space) +let uv_x = coord.x / resolution.x; // Horizontal position [0,1] +let uv_y = coord.y / resolution.y; // Vertical position [0,1] + +// Slot 6: Multi-frequency position encoding +let sin20_y = sin(20.0 * uv_y); // Periodic feature (frequency=20, vertical) + +// Slot 7: Bias dimension (always 1.0) +let bias = 1.0; // Learned bias per output channel + +// Packed storage: [p0, p1, p2, p3, uv.x, uv.y, sin(20*uv.y), 1.0] +``` + +### Input Channel Mapping + +**Weight tensor layout (12 input channels per layer):** + +| Input Channel | Feature | Description | +|--------------|---------|-------------| +| 0-3 | Previous layer output | 4D RGBA from prior CNN layer (or input RGBD for Layer 0) | +| 4-11 | Static features | 8D: p0, p1, p2, p3, uv_x, uv_y, sin20_y, bias | + +**Static feature channel details:** +- Channel 4 → p0 (RGB.r from mip level) +- Channel 5 → p1 (RGB.g from mip level) +- Channel 6 → p2 (RGB.b from mip level) +- Channel 7 → p3 (depth or RGB channel from mip level) +- Channel 8 → p4 (uv_x: normalized horizontal position) +- Channel 9 → p5 (uv_y: normalized vertical position) +- Channel 10 → p6 (sin(20*uv_y): periodic encoding) +- Channel 11 → p7 (bias: constant 1.0) + +**Note:** When generating identity weights, p4-p7 correspond to input channels 8-11, not 4-7. + +### Feature Rationale + +| Feature | Dimension | Purpose | Priority | +|---------|-----------|---------|----------| +| p0-p3 | 4D | Parametric auxiliary features (mips, gradients, etc.) | Essential | +| UV coords | 2D | Spatial position awareness | Essential | +| sin(20\*uv.y) | 1D | Periodic position encoding (vertical) | Medium | +| Bias | 1D | Learned bias (standard NN) | Essential | + +**Note:** Input image RGBD (mip 0) fed only to Layer 0. Subsequent layers see static features + previous layer output. + +**Why bias as static feature:** +- Simpler shader code (single weight array) +- Standard NN formulation: y = Wx (x includes bias term) +- Saves 56-112 bytes (no separate bias buffer) +- 7 features sufficient for initial implementation + +### Future Feature Extensions + +**Option: Additional encodings:** +- `sin(40*uv.y)` - Higher frequency encoding +- `gray_mip1` - Multi-scale luminance +- `dx`, `dy` - Sobel gradients +- `variance` - Local texture measure +- `laplacian` - Edge detection + +**Option: uint8 packing (16+ features):** +```wgsl +// texture_storage_2d stores 16 uint8 values +// Trade precision for feature count +// [R, G, B, D, uv.x, uv.y, sin10.x, sin10.y, +// sin20.x, sin20.y, dx, dy, gray_mip1, gray_mip2, var, bias] +``` +Requires quantization-aware training. + +--- + +## Layer Structure + +### Example 3-Layer Network + +``` +Layer 0: input RGBD (4D) + static (8D) = 12D → 4 channels (3×3 kernel) +Layer 1: previous (4D) + static (8D) = 12D → 4 channels (3×3 kernel) +Layer 2: previous (4D) + static (8D) = 12D → 4 channels (3×3 kernel, output RGBA) +``` + +**Output:** 4 channels (RGBA). Training targets preserve alpha from target images. + +### Weight Calculations + +**Per-layer weights (uniform 12D→4D, 3×3 kernels):** +``` +Layer 0: 12 × 3 × 3 × 4 = 432 weights +Layer 1: 12 × 3 × 3 × 4 = 432 weights +Layer 2: 12 × 3 × 3 × 4 = 432 weights +Total: 1296 weights +``` + +**Storage sizes:** +- f32: 1296 × 4 = 5,184 bytes (~5.1 KB) +- f16: 1296 × 2 = 2,592 bytes (~2.5 KB) ✓ **recommended** + +**Comparison to v1:** +- v1: ~800 weights (3.2 KB f32) +- v2: ~1296 weights (2.5 KB f16) +- **Uniform architecture, smaller than v1 f32** + +### Kernel Size Guidelines + +**1×1 kernel (pointwise):** +- No spatial context, channel mixing only +- Weights: `12 × 4 = 48` per layer +- Use for: Fast inference, channel remapping + +**3×3 kernel (standard conv):** +- Local spatial context (recommended) +- Weights: `12 × 9 × 4 = 432` per layer +- Use for: Most layers (balanced quality/size) + +**5×5 kernel (large receptive field):** +- Wide spatial context +- Weights: `12 × 25 × 4 = 1200` per layer +- Use for: Output layer, fine detail enhancement + +### Channel Storage (4×f16 per texel) + +```wgsl +@group(0) @binding(1) var layer_input: texture_2d; + +fn unpack_channels(coord: vec2) -> vec4 { + let packed = textureLoad(layer_input, coord, 0); + let v0 = unpack2x16float(packed.x); // [ch0, ch1] + let v1 = unpack2x16float(packed.y); // [ch2, ch3] + return vec4(v0.x, v0.y, v1.x, v1.y); +} + +fn pack_channels(values: vec4) -> vec4 { + return vec4( + pack2x16float(vec2(values.x, values.y)), + pack2x16float(vec2(values.z, values.w)), + 0u, // Unused + 0u // Unused + ); +} +``` + +--- + +## Training Workflow + +### Script: `training/train_cnn_v2.py` + +**Static Feature Extraction:** + +```python +def compute_static_features(rgb, depth, mip_level=0): + """Generate parametric features (8D: p0-p3 + spatial). + + Args: + mip_level: 0=original, 1=half res, 2=quarter res, 3=eighth res + """ + h, w = rgb.shape[:2] + + # Generate mip level for p0-p3 (downsample then upsample) + if mip_level > 0: + mip_rgb = rgb.copy() + for _ in range(mip_level): + mip_rgb = cv2.pyrDown(mip_rgb) + for _ in range(mip_level): + mip_rgb = cv2.pyrUp(mip_rgb) + if mip_rgb.shape[:2] != (h, w): + mip_rgb = cv2.resize(mip_rgb, (w, h), interpolation=cv2.INTER_LINEAR) + else: + mip_rgb = rgb + + # Parametric features from mip level + p0, p1, p2, p3 = mip_rgb[..., 0], mip_rgb[..., 1], mip_rgb[..., 2], depth + + # UV coordinates (normalized) + uv_x = np.linspace(0, 1, w)[None, :].repeat(h, axis=0) + uv_y = np.linspace(0, 1, h)[:, None].repeat(w, axis=1) + + # Multi-frequency position encoding + sin10_x = np.sin(10.0 * uv_x) + + # Bias dimension (always 1.0) + bias = np.ones_like(p0) + + # Stack: [p0, p1, p2, p3, uv.x, uv.y, sin10_x, bias] + return np.stack([p0, p1, p2, p3, uv_x, uv_y, sin10_x, bias], axis=-1) +``` + +**Network Definition:** + +```python +class CNNv2(nn.Module): + def __init__(self, kernel_sizes, num_layers=3): + super().__init__() + if isinstance(kernel_sizes, int): + kernel_sizes = [kernel_sizes] * num_layers + self.kernel_sizes = kernel_sizes + self.layers = nn.ModuleList() + + # All layers: 12D input (4 prev + 8 static) → 4D output + for kernel_size in kernel_sizes: + self.layers.append( + nn.Conv2d(12, 4, kernel_size=kernel_size, + padding=kernel_size//2, bias=False) + ) + + def forward(self, input_rgbd, static_features): + # Layer 0: input RGBD (4D) + static (8D) = 12D + x = torch.cat([input_rgbd, static_features], dim=1) + x = self.layers[0](x) + x = torch.sigmoid(x) # Soft [0,1] for layer 0 + + # Layer 1+: previous output (4D) + static (8D) = 12D + for i in range(1, len(self.layers)): + x_input = torch.cat([x, static_features], dim=1) + x = self.layers[i](x_input) + if i < len(self.layers) - 1: + x = F.relu(x) + else: + x = torch.sigmoid(x) # Soft [0,1] for final layer + + return x # RGBA output +``` + +**Training Configuration:** + +```python +# Hyperparameters +kernel_sizes = [3, 3, 3] # Per-layer kernel sizes (e.g., [1,3,5]) +num_layers = 3 # Number of CNN layers +mip_level = 0 # Mip level for p0-p3: 0=orig, 1=half, 2=quarter, 3=eighth +grayscale_loss = False # Compute loss on grayscale (Y) instead of RGBA +learning_rate = 1e-3 +batch_size = 16 +epochs = 5000 + +# Dataset: Input RGB, Target RGBA (preserves alpha channel from image) +# Model outputs RGBA, loss compares all 4 channels (or grayscale if --grayscale-loss) + +# Training loop (standard PyTorch f32) +for epoch in range(epochs): + for rgb_batch, depth_batch, target_batch in dataloader: + # Compute static features (8D) with mip level + static_feat = compute_static_features(rgb_batch, depth_batch, mip_level) + + # Input RGBD (4D) + input_rgbd = torch.cat([rgb_batch, depth_batch.unsqueeze(1)], dim=1) + + # Forward pass + output = model(input_rgbd, static_feat) + + # Loss computation (grayscale or RGBA) + if grayscale_loss: + # Convert RGBA to grayscale: Y = 0.299*R + 0.587*G + 0.114*B + output_gray = 0.299 * output[:, 0:1] + 0.587 * output[:, 1:2] + 0.114 * output[:, 2:3] + target_gray = 0.299 * target[:, 0:1] + 0.587 * target[:, 1:2] + 0.114 * target[:, 2:3] + loss = criterion(output_gray, target_gray) + else: + loss = criterion(output, target_batch) + + # Backward pass + optimizer.zero_grad() + loss.backward() + optimizer.step() +``` + +**Checkpoint Format:** + +```python +torch.save({ + 'state_dict': model.state_dict(), # f32 weights + 'config': { + 'kernel_sizes': [3, 3, 3], # Per-layer kernel sizes + 'num_layers': 3, + 'mip_level': 0, # Mip level used for p0-p3 + 'grayscale_loss': False, # Whether grayscale loss was used + 'features': ['p0', 'p1', 'p2', 'p3', 'uv.x', 'uv.y', 'sin10_x', 'bias'] + }, + 'epoch': epoch, + 'loss': loss.item() +}, f'checkpoints/checkpoint_epoch_{epoch}.pth') +``` + +--- + +## Export Workflow + +### Script: `training/export_cnn_v2_shader.py` + +**Process:** +1. Load checkpoint (f32 PyTorch weights) +2. Extract layer configs (kernels, channels) +3. Quantize weights to float16: `weights_f16 = weights_f32.astype(np.float16)` +4. Generate WGSL shader per layer +5. Write to `workspaces//shaders/cnn_v2/cnn_v2_*.wgsl` + +**Example Generated Shader:** + +```wgsl +// cnn_v2_layer_0.wgsl - Auto-generated from checkpoint_epoch_5000.pth + +const KERNEL_SIZE: u32 = 1u; +const IN_CHANNELS: u32 = 8u; // 7 features + bias +const OUT_CHANNELS: u32 = 16u; + +// Weights quantized to float16 (stored as f32 in shader) +const weights: array = array( + 0.123047, -0.089844, 0.234375, 0.456055, ... +); + +@group(0) @binding(0) var static_features: texture_2d; +@group(0) @binding(1) var output_texture: texture_storage_2d; + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) { + // Load static features (8D) + let static_feat = get_static_features(vec2(id.xy)); + + // Convolution (1×1 kernel = pointwise) + var output: array; + for (var c: u32 = 0u; c < OUT_CHANNELS; c++) { + var sum: f32 = 0.0; + for (var k: u32 = 0u; k < IN_CHANNELS; k++) { + sum += weights[c * IN_CHANNELS + k] * static_feat[k]; + } + output[c] = max(0.0, sum); // ReLU activation + } + + // Pack and store (8×f16 per texel) + textureStore(output_texture, vec2(id.xy), pack_f16x8(output)); +} +``` + +**Float16 Quantization:** +- Training uses f32 throughout (PyTorch standard) +- Export converts to np.float16, then back to f32 for WGSL literals +- **Expected discrepancy:** <0.1% MSE (acceptable) +- Validation via HTML tool (see below) + +--- + +## Validation Workflow + +### HTML Tool: `tools/cnn_v2_test/index.html` + +**WebGPU-based testing tool** with layer visualization. + +**Usage:** +1. Open `tools/cnn_v2_test/index.html` in browser +2. Drop `.bin` weights file (from `export_cnn_v2_weights.py`) +3. Drop PNG test image +4. View results with layer inspection + +**Features:** +- Live CNN inference with WebGPU +- Layer-by-layer visualization (static features + all CNN layers) +- Weight visualization (per-layer kernels) +- View modes: CNN output, original, diff (×10) +- Blend control for comparing with original + +**Export weights:** +```bash +./training/export_cnn_v2_weights.py checkpoints/checkpoint_epoch_100.pth \ + --output-weights workspaces/main/cnn_v2_weights.bin +``` + +See `doc/CNN_V2_WEB_TOOL.md` for detailed documentation + +--- + +## Implementation Checklist + +### Phase 1: Shaders (Core Infrastructure) + +- [ ] `workspaces/main/shaders/cnn_v2/cnn_v2_static.wgsl` - Static features compute + - [ ] RGBD sampling from framebuffer + - [ ] UV coordinate calculation + - [ ] sin(10\*uv.x) computation + - [ ] Bias dimension (constant 1.0) + - [ ] Float16 packing via `pack2x16float()` + - [ ] Output to `texture_storage_2d` + +- [ ] `workspaces/main/shaders/cnn_v2/cnn_v2_layer_template.wgsl` - Layer template + - [ ] Static features unpacking + - [ ] Previous layer unpacking (8×f16) + - [ ] Convolution implementation (1×1, 3×3, 5×5) + - [ ] ReLU activation + - [ ] Output packing (8×f16) + - [ ] Proper padding handling + +### Phase 2: C++ Effect Class + +- [ ] `src/effects/cnn_v2_effect.h` - Header + - [ ] Class declaration inheriting from `PostProcessEffect` + - [ ] Static features texture member + - [ ] Layer textures vector + - [ ] Pipeline and bind group members + +- [ ] `src/effects/cnn_v2_effect.cc` - Implementation + - [ ] Constructor: Load shaders, create textures + - [ ] `init()`: Create pipelines, bind groups + - [ ] `render()`: Multi-pass execution + - [ ] Pass 0: Compute static features + - [ ] Pass 1-N: CNN layers + - [ ] Final: Composite to output + - [ ] Proper resource cleanup + +- [ ] Integration + - [ ] Add to `src/gpu/demo_effects.h` includes + - [ ] Add `cnn_v2_effect.cc` to `CMakeLists.txt` (headless + normal) + - [ ] Add shaders to `workspaces/main/assets.txt` + - [ ] Add to `src/tests/gpu/test_demo_effects.cc` + +### Phase 3: Training Pipeline + +- [ ] `training/train_cnn_v2.py` - Training script + - [ ] Static feature extraction function + - [ ] CNNv2 PyTorch model class + - [ ] Patch-based dataloader + - [ ] Training loop with checkpointing + - [ ] Command-line argument parsing + - [ ] Inference mode (ground truth generation) + +- [ ] `training/export_cnn_v2_shader.py` - Export script + - [ ] Checkpoint loading + - [ ] Weight extraction and f16 quantization + - [ ] Per-layer WGSL generation + - [ ] File output to workspace shaders/ + - [ ] Metadata preservation + +### Phase 4: Tools & Validation + +- [x] HTML validation tool - WebGPU inference with layer visualization + - [ ] Command-line argument parsing + - [ ] Shader export orchestration + - [ ] Build orchestration + - [ ] Batch image processing + - [ ] Results display + +- [ ] `src/tools/cnn_test_main.cc` - Tool updates + - [ ] Add `--cnn-version v2` flag + - [ ] CNNv2Effect instantiation path + - [ ] Static features pass execution + - [ ] Multi-layer processing + +### Phase 5: Documentation + +- [ ] `doc/HOWTO.md` - Usage guide + - [ ] Training section (CNN v2) + - [ ] Export section + - [ ] Validation section + - [ ] Examples + +- [ ] `README.md` - Project overview update + - [ ] Mention CNN v2 capability + +--- + +## File Structure + +### New Files + +``` +# Shaders (generated by export script) +workspaces/main/shaders/cnn_v2/cnn_v2_static.wgsl # Static features compute +workspaces/main/shaders/cnn_v2/cnn_v2_layer_0.wgsl # Input layer (generated) +workspaces/main/shaders/cnn_v2/cnn_v2_layer_1.wgsl # Inner layer (generated) +workspaces/main/shaders/cnn_v2/cnn_v2_layer_2.wgsl # Output layer (generated) + +# C++ implementation +src/effects/cnn_v2_effect.h # Effect class header +src/effects/cnn_v2_effect.cc # Effect implementation + +# Python training/export +training/train_cnn_v2.py # Training script +training/export_cnn_v2_shader.py # Shader generator +training/validation/ # Test images directory + +# Validation +tools/cnn_v2_test/index.html # WebGPU validation tool + +# Documentation +doc/CNN_V2.md # This file +``` + +### Modified Files + +``` +src/gpu/demo_effects.h # Add CNNv2Effect include +CMakeLists.txt # Add cnn_v2_effect.cc +workspaces/main/assets.txt # Add cnn_v2 shaders +workspaces/main/timeline.seq # Optional: add CNNv2Effect +src/tests/gpu/test_demo_effects.cc # Add CNNv2 test case +src/tools/cnn_test_main.cc # Add --cnn-version v2 +doc/HOWTO.md # Add CNN v2 sections +TODO.md # Add CNN v2 task +``` + +### Unchanged (v1 Preserved) + +``` +training/train_cnn.py # Original training +src/effects/cnn_effect.* # Original effect +workspaces/main/shaders/cnn_*.wgsl # Original v1 shaders +``` + +--- + +## Performance Characteristics + +### Static Features Compute +- **Cost:** ~0.1ms @ 1080p +- **Frequency:** Once per frame +- **Operations:** sin(), texture sampling, packing + +### CNN Layers (Example 3-layer) +- **Layer0 (1×1, 8→16):** ~0.3ms +- **Layer1 (3×3, 23→8):** ~0.8ms +- **Layer2 (5×5, 15→4):** ~1.2ms +- **Total:** ~2.4ms @ 1080p + +### Memory Usage +- Static features: 1920×1080×8×2 = 33 MB (f16) +- Layer buffers: 1920×1080×16×2 = 66 MB (max 16 channels) +- Weights: ~6.4 KB (f16, in shader code) +- **Total GPU memory:** ~100 MB + +--- + +## Size Budget + +### CNN v1 vs v2 + +| Metric | v1 | v2 | Delta | +|--------|----|----|-------| +| Weights (count) | 800 | 3268 | +2468 | +| Storage (f32) | 3.2 KB | 13.1 KB | +9.9 KB | +| Storage (f16) | N/A | 6.5 KB | +6.5 KB | +| Shader code | ~500 lines | ~800 lines | +300 lines | + +### Mitigation Strategies + +**Reduce channels:** +- [16,8,4] → [8,4,4] saves ~50% weights +- [16,8,4] → [4,4,4] saves ~60% weights + +**Smaller kernels:** +- [1,3,5] → [1,3,3] saves ~30% weights +- [1,3,5] → [1,1,3] saves ~50% weights + +**Quantization:** +- int8 weights: saves 75% (requires QAT training) +- 4-bit weights: saves 87.5% (extreme, needs research) + +**Target:** Keep CNN v2 under 10 KB for 64k demo constraint + +--- + +## Future Extensions + +### Flexible Feature Layout (Binary Format v3) + +**TODO:** Support arbitrary feature vector layouts and ordering in binary format. + +**Current Limitation:** +- Feature layout hardcoded: `[p0, p1, p2, p3, uv_x, uv_y, sin10_x, bias]` +- Shader must match training script exactly +- Experimentation requires shader recompilation + +**Proposed Enhancement:** +- Add feature descriptor to binary format header +- Specify feature types, sources, and ordering +- Runtime shader generation or dynamic feature indexing +- Examples: `[R, G, B, dx, dy, uv_x, bias]` or `[mip1.r, mip2.g, laplacian, uv_x, sin20_x, bias]` + +**Benefits:** +- Training experiments without C++/shader changes +- A/B test different feature combinations +- Single binary format, multiple architectures +- Faster iteration on feature engineering + +**Implementation Options:** +1. **Static approach:** Generate shader code from descriptor at load time +2. **Dynamic approach:** Array-based indexing with feature map uniform +3. **Hybrid:** Precompile common layouts, fallback to dynamic + +See `doc/CNN_V2_BINARY_FORMAT.md` for proposed descriptor format. + +--- + +### More Features (uint8 Packing) + +```wgsl +// 16 uint8 features per texel (texture_storage_2d) +// [R, G, B, D, uv.x, uv.y, sin10.x, sin10.y, +// sin20.x, sin20.y, dx, dy, gray_mip1, gray_mip2, variance, bias] +``` +- Trade precision for quantity +- Requires quantization-aware training + +### Temporal Features + +- Previous frame RGBA (motion awareness) +- Optical flow vectors +- Requires multi-frame buffer + +### Learned Position Encodings + +- Replace hand-crafted sin(10\*uv) with learned embeddings +- Requires separate embedding network +- Similar to NeRF position encoding + +### Dynamic Architecture + +- Runtime kernel size selection based on scene +- Conditional layer execution (skip connections) +- Layer pruning for performance + +--- + +## References + +- **v1 Implementation:** `src/effects/cnn_effect.*` +- **Training Guide:** `doc/HOWTO.md` (CNN Training section) +- **Test Tool:** `doc/CNN_TEST_TOOL.md` +- **Shader System:** `doc/SEQUENCE.md` +- **Size Measurement:** `doc/SIZE_MEASUREMENT.md` + +--- + +## Appendix: Design Decisions + +### Why Bias as Static Feature? + +**Alternatives considered:** +1. Separate bias array per layer (Option B) +2. Bias as static feature = 1.0 (Option A, chosen) + +**Decision rationale:** +- Simpler shader code (fewer bindings) +- Standard NN formulation (augmented input) +- Saves 56-112 bytes per model +- 7 features sufficient for v1 implementation +- Can extend to uint8 packing if >7 features needed + +### Why Float16 for Weights? + +**Alternatives considered:** +1. Keep f32 (larger, more accurate) +2. Use f16 (smaller, GPU-native) +3. Use int8 (smallest, needs QAT) + +**Decision rationale:** +- f16 saves 50% vs f32 (critical for 64k target) +- GPU-native support (pack2x16float in WGSL) +- <0.1% accuracy loss (acceptable) +- Simpler than int8 quantization + +### Why Multi-Frequency Position Encoding? + +**Inspiration:** NeRF (Neural Radiance Fields) + +**Benefits:** +- Helps network learn high-frequency details +- Better than raw UV coordinates +- Small footprint (1D per frequency) + +**Future:** Add sin(20\*uv), sin(40\*uv) if >7 features available + +--- + +## Related Documentation + +- `doc/CNN_V2_BINARY_FORMAT.md` - Binary weight file specification (.bin format) +- `doc/CNN_V2_WEB_TOOL.md` - WebGPU testing tool with layer visualization +- `doc/CNN_TEST_TOOL.md` - C++ offline validation tool (deprecated) +- `doc/HOWTO.md` - Training and validation workflows + +--- + +**Document Version:** 1.0 +**Last Updated:** 2026-02-12 +**Status:** Design approved, ready for implementation diff --git a/cnn_v2/docs/CNN_V2_BINARY_FORMAT.md b/cnn_v2/docs/CNN_V2_BINARY_FORMAT.md new file mode 100644 index 0000000..59c859d --- /dev/null +++ b/cnn_v2/docs/CNN_V2_BINARY_FORMAT.md @@ -0,0 +1,235 @@ +# CNN v2 Binary Weight Format Specification + +Binary format for storing trained CNN v2 weights with static feature architecture. + +**File Extension:** `.bin` +**Byte Order:** Little-endian +**Version:** 2.0 (supports mip-level for parametric features) +**Backward Compatible:** Version 1.0 files supported (mip_level=0) + +--- + +## File Structure + +**Version 2 (current):** +``` +┌─────────────────────┐ +│ Header (20 bytes) │ +├─────────────────────┤ +│ Layer Info │ +│ (20 bytes × N) │ +├─────────────────────┤ +│ Weight Data │ +│ (variable size) │ +└─────────────────────┘ +``` + +**Version 1 (legacy):** +``` +┌─────────────────────┐ +│ Header (16 bytes) │ +├─────────────────────┤ +│ Layer Info │ +│ (20 bytes × N) │ +├─────────────────────┤ +│ Weight Data │ +│ (variable size) │ +└─────────────────────┘ +``` + +--- + +## Header + +**Version 2 (20 bytes):** + +| Offset | Type | Field | Description | +|--------|------|----------------|--------------------------------------| +| 0x00 | u32 | magic | Magic number: `0x32_4E_4E_43` ("CNN2") | +| 0x04 | u32 | version | Format version (2 for current) | +| 0x08 | u32 | num_layers | Number of CNN layers (excludes static features) | +| 0x0C | u32 | total_weights | Total f16 weight count across all layers | +| 0x10 | u32 | mip_level | Mip level for p0-p3 features (0=original, 1=half, 2=quarter, 3=eighth) | + +**Version 1 (16 bytes) - Legacy:** + +| Offset | Type | Field | Description | +|--------|------|----------------|--------------------------------------| +| 0x00 | u32 | magic | Magic number: `0x32_4E_4E_43` ("CNN2") | +| 0x04 | u32 | version | Format version (1) | +| 0x08 | u32 | num_layers | Number of CNN layers | +| 0x0C | u32 | total_weights | Total f16 weight count | + +**Note:** Loaders should check version field and handle both formats. Version 1 files treated as mip_level=0. + +--- + +## Layer Info (20 bytes per layer) + +Repeated `num_layers` times: +- **Version 2:** Starting at offset 0x14 (20 bytes) +- **Version 1:** Starting at offset 0x10 (16 bytes) + +| Offset | Type | Field | Description | +|-------------|------|----------------|--------------------------------------| +| 0x00 | u32 | kernel_size | Convolution kernel dimension (3, 5, 7, etc.) | +| 0x04 | u32 | in_channels | Input channel count (includes 8 static features for Layer 1) | +| 0x08 | u32 | out_channels | Output channel count (max 8) | +| 0x0C | u32 | weight_offset | Weight array start index (f16 units, relative to weight data section) | +| 0x10 | u32 | weight_count | Number of f16 weights for this layer | + +**Layer Order:** Sequential (Layer 1, Layer 2, Layer 3, ...) + +--- + +## Weight Data (variable size) + +Starts at offset: +- **Version 2:** `20 + (num_layers × 20)` +- **Version 1:** `16 + (num_layers × 20)` + +**Format:** Packed f16 pairs stored as u32 +**Packing:** `u32 = (f16_hi << 16) | f16_lo` +**Storage:** Sequential by layer, then by output channel, input channel, spatial position + +**Weight Indexing:** +``` +weight_idx = output_ch × (in_channels × kernel_size²) + + input_ch × kernel_size² + + (ky × kernel_size + kx) +``` + +Where: +- `output_ch` ∈ [0, out_channels) +- `input_ch` ∈ [0, in_channels) +- `ky`, `kx` ∈ [0, kernel_size) + +**Unpacking f16 from u32:** +```c +uint32_t packed = weights_buffer[weight_idx / 2]; +uint16_t f16_bits = (weight_idx % 2 == 0) ? (packed & 0xFFFF) : (packed >> 16); +``` + +--- + +## Example: 3-Layer Network (Version 2) + +**Configuration:** +- Mip level: 0 (original resolution) +- Layer 0: 12→4, kernel 3×3 (432 weights) +- Layer 1: 12→4, kernel 3×3 (432 weights) +- Layer 2: 12→4, kernel 3×3 (432 weights) + +**File Layout:** +``` +Offset Size Content +------ ---- ------- +0x00 20 Header (magic, version=2, layers=3, weights=1296, mip_level=0) +0x14 20 Layer 0 info (kernel=3, in=12, out=4, offset=0, count=432) +0x28 20 Layer 1 info (kernel=3, in=12, out=4, offset=432, count=432) +0x3C 20 Layer 2 info (kernel=3, in=12, out=4, offset=864, count=432) +0x50 2592 Weight data (1296 u32 packed f16 pairs) + ---- +Total: 2672 bytes (~2.6 KB) +``` + +--- + +## Static Features + +Not stored in .bin file (computed at runtime): + +**8D Input Features:** +1. **p0** - Parametric feature 0 (from mip level) +2. **p1** - Parametric feature 1 (from mip level) +3. **p2** - Parametric feature 2 (from mip level) +4. **p3** - Parametric feature 3 (depth or from mip level) +5. **UV_X** - Normalized x coordinate [0,1] +6. **UV_Y** - Normalized y coordinate [0,1] +7. **sin(20 × UV_Y)** - Spatial frequency encoding (vertical, frequency=20) +8. **1.0** - Bias term + +**Mip Level Usage (p0-p3):** +- `mip_level=0`: RGB from original resolution (mip 0) +- `mip_level=1`: RGB from half resolution (mip 1), upsampled +- `mip_level=2`: RGB from quarter resolution (mip 2), upsampled +- `mip_level=3`: RGB from eighth resolution (mip 3), upsampled + +**Layer 0** receives input RGBD (4D) + static features (8D) = 12D input → 4D output. +**Layer 1+** receive previous layer output (4D) + static features (8D) = 12D input → 4D output. + +--- + +## Validation + +**Magic Check:** +```c +uint32_t magic; +fread(&magic, 4, 1, fp); +if (magic != 0x32_4E_4E_43) { error("Invalid CNN v2 file"); } +``` + +**Version Check:** +```c +uint32_t version; +fread(&version, 4, 1, fp); +if (version != 1 && version != 2) { error("Unsupported version"); } +uint32_t header_size = (version == 1) ? 16 : 20; +``` + +**Size Check:** +```c +expected_size = header_size + (num_layers × 20) + (total_weights × 2); +if (file_size != expected_size) { error("Size mismatch"); } +``` + +**Weight Offset Sanity:** +```c +// Each layer's offset should match cumulative count +uint32_t cumulative = 0; +for (int i = 0; i < num_layers; i++) { + if (layers[i].weight_offset != cumulative) { error("Invalid offset"); } + cumulative += layers[i].weight_count; +} +if (cumulative != total_weights) { error("Total mismatch"); } +``` + +--- + +## Future Extensions + +**TODO: Flexible Feature Layout** + +Current limitation: Feature vector layout is hardcoded as `[p0, p1, p2, p3, uv_x, uv_y, sin10_x, bias]`. + +Proposed enhancement for version 3: +- Add feature descriptor section to header +- Specify feature count, types, and ordering +- Support arbitrary 7D feature combinations (e.g., `[R, G, B, dx, dy, uv_x, bias]`) +- Allow runtime shader generation based on descriptor +- Enable experimentation without recompiling shaders + +Example descriptor format: +``` +struct FeatureDescriptor { + u32 feature_count; // Number of features (typically 7-8) + u32 feature_types[8]; // Type enum per feature + u32 feature_sources[8]; // Source enum (mip0, mip1, gradient, etc.) + u32 reserved[8]; // Future use +} +``` + +Benefits: +- Training can experiment with different feature combinations +- No shader recompilation needed +- Single binary format supports multiple architectures +- Easier A/B testing of feature effectiveness + +--- + +## Related Files + +- `training/export_cnn_v2_weights.py` - Binary export tool +- `src/effects/cnn_v2_effect.cc` - C++ loader +- `tools/cnn_v2_test/index.html` - WebGPU validator +- `doc/CNN_V2.md` - Architecture design diff --git a/cnn_v2/docs/CNN_V2_DEBUG_TOOLS.md b/cnn_v2/docs/CNN_V2_DEBUG_TOOLS.md new file mode 100644 index 0000000..8d1289a --- /dev/null +++ b/cnn_v2/docs/CNN_V2_DEBUG_TOOLS.md @@ -0,0 +1,143 @@ +# CNN v2 Debugging Tools + +Tools for investigating CNN v2 mismatch between HTML tool and cnn_test. + +--- + +## Identity Weight Generator + +**Purpose:** Generate trivial .bin files with identity passthrough for debugging. + +**Script:** `training/gen_identity_weights.py` + +**Usage:** +```bash +# 1×1 identity (default) +./training/gen_identity_weights.py workspaces/main/weights/cnn_v2_identity.bin + +# 3×3 identity +./training/gen_identity_weights.py workspaces/main/weights/cnn_v2_identity_3x3.bin --kernel-size 3 + +# Mix mode: 50-50 blend (0.5*p0+0.5*p4, etc) +./training/gen_identity_weights.py output.bin --mix + +# Static features only: p4→ch0, p5→ch1, p6→ch2, p7→ch3 +./training/gen_identity_weights.py output.bin --p47 + +# Custom mip level +./training/gen_identity_weights.py output.bin --kernel-size 1 --mip-level 2 +``` + +**Output:** +- Single layer, 12D→4D (4 input channels + 8 static features) +- Identity mode: Output Ch{0,1,2,3} = Input Ch{0,1,2,3} +- Mix mode (--mix): Output Ch{i} = 0.5*Input Ch{i} + 0.5*Input Ch{i+4} (50-50 blend, avoids overflow) +- Static mode (--p47): Output Ch{i} = Input Ch{i+4} (static features only, visualizes p4-p7) +- Minimal file size (~136 bytes for 1×1, ~904 bytes for 3×3) + +**Validation:** +Load in HTML tool or cnn_test - output should match input (RGB only, ignoring static features). + +--- + +## Composited Layer Visualization + +**Purpose:** Save current layer view as single composited image (4 channels side-by-side, grayscale). + +**Location:** HTML tool - "Layer Visualization" panel + +**Usage:** +1. Load image + weights in HTML tool +2. Select layer to visualize (Static 0-3, Static 4-7, Layer 0, Layer 1, etc.) +3. Click "Save Composited" button +4. Downloads PNG: `composited_layer{N}_{W}x{H}.png` + +**Output:** +- 4 channels stacked horizontally +- Grayscale representation +- Useful for comparing layer activations across tools + +--- + +## Debugging Strategy + +### Track a) Binary Conversion Chain + +**Hypothesis:** Conversion error in .bin ↔ base64 ↔ Float32Array + +**Test:** +1. Generate identity weights: + ```bash + ./training/gen_identity_weights.py workspaces/main/weights/test_identity.bin + ``` + +2. Load in HTML tool - output should match input RGB + +3. If mismatch: + - Check Python export: f16 packing in `export_cnn_v2_weights.py` line 105 + - Check HTML parsing: `unpackF16()` in `index.html` line 805-815 + - Check weight indexing: `get_weight()` shader function + +**Key locations:** +- Python: `np.float16` → `view(np.uint32)` (line 105 of export script) +- JS: `DataView` → `unpackF16()` → manual f16 decode (line 773-803) +- WGSL: `unpack2x16float()` built-in (line 492 of shader) + +### Track b) Layer Visualization + +**Purpose:** Confirm layer outputs match between HTML and C++ + +**Method:** +1. Run identical input through both tools +2. Save composited layers from HTML tool +3. Compare with cnn_test output +4. Use identity weights to isolate weight loading from computation + +### Track c) Trivial Test Case + +**Use identity weights to test:** +- Weight loading (binary parsing) +- Feature generation (static features) +- Convolution (should be passthrough) +- Output packing + +**Expected behavior:** +- Input RGB → Output RGB (exact match) +- Static features ignored (all zeros in identity matrix) + +--- + +## Known Issues + +### ~~Layer 0 Visualization Scale~~ [FIXED] + +**Issue:** Layer 0 output displayed at 0.5× brightness (divided by 2). + +**Cause:** Line 1530 used `vizScale = 0.5` for all CNN layers, but Layer 0 is clamped [0,1] and doesn't need dimming. + +**Fix:** Use scale 1.0 for Layer 0 output (layerIdx=1), 0.5 only for middle layers (ReLU, unbounded). + +### Remaining Mismatch + +**Current:** HTML tool and cnn_test produce different outputs for same input/weights. + +**Suspects:** +1. F16 unpacking difference (CPU vs GPU vs JS) +2. Static feature generation (RGBD, UV, sin encoding) +3. Convolution kernel iteration order +4. Output packing/unpacking + +**Next steps:** +1. Test with identity weights (eliminates weight loading) +2. Compare composited layer outputs +3. Add debug visualization for static features +4. Hex dump comparison (first 8 pixels) - use `--debug-hex` flag in cnn_test + +--- + +## Related Documentation + +- `doc/CNN_V2.md` - CNN v2 architecture +- `doc/CNN_V2_WEB_TOOL.md` - HTML tool documentation +- `doc/CNN_TEST_TOOL.md` - cnn_test CLI tool +- `training/export_cnn_v2_weights.py` - Binary export format diff --git a/cnn_v2/docs/CNN_V2_WEB_TOOL.md b/cnn_v2/docs/CNN_V2_WEB_TOOL.md new file mode 100644 index 0000000..b6f5b0b --- /dev/null +++ b/cnn_v2/docs/CNN_V2_WEB_TOOL.md @@ -0,0 +1,348 @@ +# CNN v2 Web Testing Tool + +Browser-based WebGPU tool for validating CNN v2 inference with layer visualization and weight inspection. + +**Location:** `tools/cnn_v2_test/index.html` + +--- + +## Status (2026-02-13) + +**Working:** +- ✅ WebGPU initialization and device setup +- ✅ Binary weight file parsing (v1 and v2 formats) +- ✅ Automatic mip-level detection from binary format v2 +- ✅ Weight statistics (min/max per layer) +- ✅ UI layout with collapsible panels +- ✅ Mode switching (Activations/Weights tabs) +- ✅ Canvas context management (2D for weights, WebGPU for activations) +- ✅ Weight visualization infrastructure (layer selection, grid layout) +- ✅ Layer naming matches codebase convention (Layer 0, Layer 1, Layer 2) +- ✅ Static features split visualization (Static 0-3, Static 4-7) +- ✅ All layers visible including output layer (Layer 2) +- ✅ Video playback support (MP4, WebM) with frame-by-frame controls +- ✅ Video looping (automatic continuous playback) +- ✅ Mip level selection (p0-p3 features at different resolutions) + +**Recent Changes (Latest):** +- Binary format v2 support: Reads mip_level from 20-byte header +- Backward compatible: v1 (16-byte header) → mip_level=0 +- Auto-update UI dropdown when loading weights with mip_level +- Display mip_level in metadata panel +- Code refactoring: Extracted FULLSCREEN_QUAD_VS shader (reused 3× across pipelines) +- Added helper methods: `getDimensions()`, `setVideoControlsEnabled()` +- Improved code organization with section headers and comments +- Moved Mip Level selector to bottom of left sidebar (removed "Features (p0-p3)" label) +- Added `loop` attribute to video element for automatic continuous playback + +**Previous Fixes:** +- Fixed Layer 2 not appearing (was excluded from layerOutputs due to isOutput check) +- Fixed canvas context switching (force clear before recreation) +- Added Static 0-3 / Static 4-7 buttons to view all 8 static feature channels +- Aligned naming with train_cnn_v2.py/.wgsl: Layer 0, Layer 1, Layer 2 (not Layer 1, 2, 3) +- Disabled Static buttons in weights mode (no learnable weights) + +**Known Issues:** +- Layer activation visualization may show black if texture data not properly unpacked +- Weight kernel display depends on correct 2D context creation after canvas recreation + +--- + +## Architecture + +### File Structure +- Single-file HTML tool (~1100 lines) +- Embedded shaders: STATIC_SHADER, CNN_SHADER, DISPLAY_SHADER, LAYER_VIZ_SHADER +- Shared WGSL component: FULLSCREEN_QUAD_VS (reused across render pipelines) +- **Embedded default weights:** DEFAULT_WEIGHTS_B64 (base64-encoded binary v2) + - Current: 4 layers (3×3, 5×5, 3×3, 3×3), 2496 f16 weights, mip_level=2 + - Source: `workspaces/main/weights/cnn_v2_weights.bin` + - Updates: Re-encode binary with `base64 -i ` and update constant +- Pure WebGPU (no external dependencies) + +### Code Organization + +**Recent Refactoring (2026-02-13):** +- Extracted `FULLSCREEN_QUAD_VS` constant: Reused fullscreen quad vertex shader (2 triangles covering NDC) +- Added helper methods to CNNTester class: + - `getDimensions()`: Returns current source dimensions (video or image) + - `setVideoControlsEnabled(enabled)`: Centralized video control enable/disable +- Consolidated duplicate vertex shader code (used in mipmap generation, display, layer visualization) +- Added section headers in JavaScript for better navigation +- Improved inline comments explaining shader architecture + +**Benefits:** +- Reduced code duplication (~40 lines saved) +- Easier maintenance (single source of truth for fullscreen quad) +- Clearer separation of concerns + +### Key Components + +**1. Weight Parsing** +- Reads binary format v2: header (20B) + layer info (20B×N) + f16 weights +- Backward compatible with v1: header (16B), mip_level defaults to 0 +- Computes min/max per layer via f16 unpacking +- Stores `{ layers[], weights[], mipLevel, fileSize }` +- Auto-sets UI mip-level dropdown from loaded weights + +**2. CNN Pipeline** +- Static features computation (RGBD + UV + sin + bias → 7D packed) +- Layer-by-layer convolution with storage buffer weights +- Ping-pong buffers for intermediate results +- Copy to persistent textures for visualization + +**3. Visualization Modes** + +**Activations Mode:** +- 4 grayscale views per layer (channels 0-3 of up to 8 total) +- WebGPU compute → unpack f16 → scale → grayscale +- Auto-scale: Static features = 1.0, CNN layers = 0.2 +- Static features: Shows R,G,B,D (first 4 of 8: RGBD+UV+sin+bias) +- CNN layers: Shows first 4 output channels + +**Weights Mode:** +- 2D canvas rendering per output channel +- Shows all input kernels horizontally +- Normalized by layer min/max → [0, 1] → grayscale +- 20px cells, 2px padding between kernels + +### Texture Management + +**Persistent Storage (layerTextures[]):** +- One texture per layer output (static + all CNN layers) +- `rgba32uint` format (packed f16 data) +- `COPY_DST` usage for storing results + +**Compute Buffers (computeTextures[]):** +- 2 textures for ping-pong computation +- Reused across all layers +- `COPY_SRC` usage for copying to persistent storage + +**Pipeline:** +``` +Static pass → copy to layerTextures[0] +For each CNN layer i: + Compute (ping-pong) → copy to layerTextures[i+1] +``` + +### Layer Indexing + +**UI Layer Buttons:** +- "Static" → layerOutputs[0] (7D input features) +- "Layer 1" → layerOutputs[1] (CNN layer 1 output, uses weights.layers[0]) +- "Layer 2" → layerOutputs[2] (CNN layer 2 output, uses weights.layers[1]) +- "Layer N" → layerOutputs[N] (CNN layer N output, uses weights.layers[N-1]) + +**Weights Table:** +- "Layer 1" → weights.layers[0] (first CNN layer weights) +- "Layer 2" → weights.layers[1] (second CNN layer weights) +- "Layer N" → weights.layers[N-1] + +**Consistency:** Both UI and weights table use same numbering (1, 2, 3...) for CNN layers. + +--- + +## Known Issues + +### Issue #1: Layer Activations Show Black + +**Symptom:** +- All 4 channel canvases render black +- UV gradient test (debug mode 10) works +- Raw packed data test (mode 11) shows black +- Unpacked f16 test (mode 12) shows black + +**Diagnosis:** +- Texture access works (UV gradient visible) +- Texture data is all zeros (packed.x = 0) +- Textures being read are empty + +**Root Cause:** +- `copyTextureToTexture` operations may not be executing +- Possible ordering issue (copies not submitted before visualization) +- Alternative: textures created with wrong usage flags + +**Investigation Steps Taken:** +1. Added `onSubmittedWorkDone()` wait before visualization +2. Verified texture creation with `COPY_SRC` and `COPY_DST` flags +3. Confirmed separate texture allocation per layer (no aliasing) +4. Added debug shader modes to isolate issue + +**Next Steps:** +- Verify encoder contains copy commands (add debug logging) +- Check if compute passes actually write data (add known-value test) +- Test copyTextureToTexture in isolation +- Consider CPU readback to verify texture contents + +### Issue #2: Weight Visualization Empty + +**Symptom:** +- Canvases created with correct dimensions (logged) +- No visual output (black canvases) +- Console logs show method execution + +**Potential Causes:** +1. Weight indexing calculation incorrect +2. Canvas not properly attached to DOM when rendering +3. 2D context operations not flushing +4. Min/max normalization producing black (all values equal?) + +**Debug Added:** +- Comprehensive logging of dimensions, indices, ranges +- Canvas context check before rendering + +**Next Steps:** +- Add test rendering (fixed gradient) to verify 2D context works +- Log sample weight values to verify data access +- Check if canvas is visible in DOM inspector +- Verify min/max calculation produces valid range + +--- + +## UI Layout + +### Header +- Controls: Blend slider, Depth input, View mode display +- Drop zone for .bin weight files + +### Content Area + +**Left Sidebar (300px):** +1. Drop zone for .bin weight files +2. Weights Info panel (file size, layer table with min/max) +3. Weights Visualization panel (per-layer kernel display) +4. **Mip Level selector** (bottom) - Select p0/p1/p2 for static features + +**Main Canvas (center):** +- CNN output display with video controls (Play/Pause, Frame ◄/►) +- Supports both PNG images and video files (MP4, WebM) +- Video loops automatically for continuous playback + +**Right Sidebar (panels):** +1. **Layer Visualization Panel** (top, flex: 1) + - Layer selection buttons (Static 0-3, Static 4-7, Layer 0, Layer 1, ...) + - 2×2 grid of channel views (grayscale activations) + - 4× zoom view at bottom + +### Footer +- Status line (GPU timing, dimensions, mode) +- Console log (scrollable, color-coded) + +--- + +## Shader Details + +### LAYER_VIZ_SHADER + +**Purpose:** Display single channel from packed layer texture + +**Inputs:** +- `@binding(0) layer_tex: texture_2d` - Packed f16 layer data +- `@binding(1) viz_params: vec2` - (channel_idx, scale) + +**Debug Modes:** +- Channel 10: UV gradient (texture coordinate test) +- Channel 11: Raw packed u32 data +- Channel 12: First unpacked f16 value + +**Normal Operation:** +- Unpack all 8 f16 channels from rgba32uint +- Select channel by index (0-7) +- Apply scale factor (1.0 for static, 0.2 for CNN) +- Clamp to [0, 1] and output grayscale + +**Scale Rationale:** +- Static features (RGBD, UV): already in [0, 1] range +- CNN activations: post-ReLU [0, ~5], need scaling for visibility + +--- + +## Binary Weight Format + +See `doc/CNN_V2_BINARY_FORMAT.md` for complete specification. + +**Quick Summary:** +- Header: 16 bytes (magic, version, layer count, total weights) +- Layer info: 20 bytes × N (kernel size, channels, offsets) +- Weights: Packed f16 pairs as u32 + +--- + +## Testing Workflow + +### Load & Parse +1. Drop PNG image → displays original +2. Drop .bin weights → parses and shows info table +3. Auto-runs CNN pipeline + +### Verify Pipeline +1. Check console for "Running CNN pipeline" +2. Verify "Completed in Xms" +3. Check "Layer visualization ready: N layers" + +### Debug Activations +1. Select "Activations" tab +2. Click layer buttons to switch +3. Check console for texture/canvas logs +4. If black: note which debug modes work (UV vs data) + +### Debug Weights +1. Select "Weights" tab +2. Click Layer 1 or Layer 2 (Layer 0 has no weights) +3. Check console for "Visualizing Layer N weights" +4. Check canvas dimensions logged +5. Verify weight range is non-trivial (not [0, 0]) + +--- + +## Integration with Main Project + +**Training Pipeline:** +```bash +# Generate weights +./training/train_cnn_v2.py --export-binary + +# Test in browser +open tools/cnn_v2_test/index.html +# Drop: workspaces/main/cnn_v2_weights.bin +# Drop: training/input/test.png +``` + +**Validation:** +- Compare against demo CNNv2Effect (visual check) +- Verify layer count matches binary file +- Check weight ranges match training logs + +--- + +## Future Enhancements + +- [ ] Fix layer activation visualization (black texture issue) +- [ ] Fix weight kernel display (empty canvas issue) +- [ ] Add per-channel auto-scaling (compute min/max from visible data) +- [ ] Export rendered outputs (download PNG) +- [ ] Side-by-side comparison with original +- [ ] Heatmap mode (color-coded activations) +- [ ] Weight statistics overlay (mean, std, sparsity) +- [ ] Batch processing (multiple images in sequence) +- [ ] Integration with Python training (live reload) + +--- + +## Code Metrics + +- Total lines: ~1100 +- JavaScript: ~700 lines +- WGSL shaders: ~300 lines +- HTML/CSS: ~100 lines + +**Dependencies:** None (pure WebGPU + HTML5) + +--- + +## Related Files + +- `doc/CNN_V2.md` - CNN v2 architecture and design +- `doc/CNN_TEST_TOOL.md` - C++ offline testing tool (deprecated) +- `training/train_cnn_v2.py` - Training script with binary export +- `workspaces/main/cnn_v2_weights.bin` - Trained weights diff --git a/cnn_v2/scripts/train_cnn_v2_full.sh b/cnn_v2/scripts/train_cnn_v2_full.sh new file mode 100755 index 0000000..a21c1ac --- /dev/null +++ b/cnn_v2/scripts/train_cnn_v2_full.sh @@ -0,0 +1,428 @@ +#!/bin/bash +# Complete CNN v2 Training Pipeline +# Train → Export → Build → Validate +# Usage: ./train_cnn_v2_full.sh [OPTIONS] +# +# MODES: +# (none) Run complete pipeline: train → export → build → validate +# --validate Validate only (skip training, use existing weights) +# --validate CHECKPOINT Validate with specific checkpoint file +# --export-only CHECKPOINT Export weights only (skip training, build, validation) +# +# TRAINING PARAMETERS: +# --epochs N Training epochs (default: 200) +# --batch-size N Batch size (default: 16) +# --lr FLOAT Learning rate (default: 1e-3) +# --checkpoint-every N Checkpoint interval (default: 50) +# --kernel-sizes K Comma-separated kernel sizes (default: 3,3,3) +# --num-layers N Number of layers (default: 3) +# --mip-level N Mip level for p0-p3 features: 0-3 (default: 0) +# --grayscale-loss Compute loss on grayscale instead of RGBA +# +# PATCH PARAMETERS: +# --patch-size N Patch size (default: 8) +# --patches-per-image N Patches per image (default: 256) +# --detector TYPE Detector: harris|fast|shi-tomasi|gradient (default: harris) +# --full-image Use full-image training (disables patch mode) +# --image-size N Image size for full-image mode (default: 256) +# +# DIRECTORIES: +# --input DIR Input directory (default: training/input) +# --target DIR Target directory (default: training/target_1) +# --checkpoint-dir DIR Checkpoint directory (default: checkpoints) +# --validation-dir DIR Validation directory (default: validation_results) +# +# OUTPUT: +# --output-weights PATH Output binary weights file (default: workspaces/main/weights/cnn_v2_weights.bin) +# +# OTHER: +# --help Show this help message +# +# Examples: +# ./train_cnn_v2_full.sh +# ./train_cnn_v2_full.sh --epochs 500 --batch-size 32 +# ./train_cnn_v2_full.sh --validate +# ./train_cnn_v2_full.sh --validate checkpoints/checkpoint_epoch_50.pth +# ./train_cnn_v2_full.sh --export-only checkpoints/checkpoint_epoch_100.pth +# ./train_cnn_v2_full.sh --mip-level 1 --kernel-sizes 3,5,3 + +set -e + +PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$PROJECT_ROOT" + +# Helper functions +export_weights() { + python3 "$SCRIPT_DIR/../training/export_cnn_v2_weights.py" "$1" --output-weights "$2" --quiet +} + +find_latest_checkpoint() { + ls -t "$CHECKPOINT_DIR"/checkpoint_epoch_*.pth 2>/dev/null | head -1 +} + +build_target() { + cmake --build build -j4 --target "$1" > /dev/null 2>&1 +} + +# Path resolution for running from any directory +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +# Default configuration +INPUT_DIR="training/input" +TARGET_DIR="training/target_1" +CHECKPOINT_DIR="checkpoints" +VALIDATION_DIR="validation_results" +EPOCHS=200 +CHECKPOINT_EVERY=50 +BATCH_SIZE=16 +LEARNING_RATE=1e-3 +PATCH_SIZE=8 +PATCHES_PER_IMAGE=256 +DETECTOR="harris" +KERNEL_SIZES="3,3,3" +NUM_LAYERS=3 +MIP_LEVEL=0 +GRAYSCALE_LOSS=false +FULL_IMAGE_MODE=false +IMAGE_SIZE=256 +OUTPUT_WEIGHTS="${PROJECT_ROOT}/workspaces/main/weights/cnn_v2_weights.bin" + +# Parse arguments +VALIDATE_ONLY=false +VALIDATE_CHECKPOINT="" +EXPORT_ONLY=false +EXPORT_CHECKPOINT="" + +if [ "$1" = "--help" ] || [ "$1" = "-h" ]; then + head -47 "$0" | grep "^#" | grep -v "^#!/" | sed 's/^# *//' + exit 0 +fi + +# Parse all arguments +while [[ $# -gt 0 ]]; do + case "$1" in + --export-only) + EXPORT_ONLY=true + if [ -z "$2" ]; then + echo "Error: --export-only requires a checkpoint file argument" + exit 1 + fi + EXPORT_CHECKPOINT="$2" + shift 2 + ;; + --validate) + VALIDATE_ONLY=true + if [ -n "$2" ] && [[ ! "$2" =~ ^-- ]]; then + VALIDATE_CHECKPOINT="$2" + shift 2 + else + shift + fi + ;; + --epochs) + if [ -z "$2" ]; then + echo "Error: --epochs requires a number argument" + exit 1 + fi + EPOCHS="$2" + shift 2 + ;; + --batch-size) + if [ -z "$2" ]; then + echo "Error: --batch-size requires a number argument" + exit 1 + fi + BATCH_SIZE="$2" + shift 2 + ;; + --checkpoint-every) + if [ -z "$2" ]; then + echo "Error: --checkpoint-every requires a number argument" + exit 1 + fi + CHECKPOINT_EVERY="$2" + shift 2 + ;; + --kernel-sizes) + if [ -z "$2" ]; then + echo "Error: --kernel-sizes requires a comma-separated list" + exit 1 + fi + KERNEL_SIZES="$2" + shift 2 + ;; + --num-layers) + if [ -z "$2" ]; then + echo "Error: --num-layers requires a number argument" + exit 1 + fi + NUM_LAYERS="$2" + shift 2 + ;; + --mip-level) + if [ -z "$2" ]; then + echo "Error: --mip-level requires a level argument (0-3)" + exit 1 + fi + MIP_LEVEL="$2" + shift 2 + ;; + --grayscale-loss) + GRAYSCALE_LOSS=true + shift + ;; + --lr) + if [ -z "$2" ]; then + echo "Error: --lr requires a float argument" + exit 1 + fi + LEARNING_RATE="$2" + shift 2 + ;; + --patch-size) + if [ -z "$2" ]; then + echo "Error: --patch-size requires a number argument" + exit 1 + fi + PATCH_SIZE="$2" + shift 2 + ;; + --patches-per-image) + if [ -z "$2" ]; then + echo "Error: --patches-per-image requires a number argument" + exit 1 + fi + PATCHES_PER_IMAGE="$2" + shift 2 + ;; + --detector) + if [ -z "$2" ]; then + echo "Error: --detector requires a type argument" + exit 1 + fi + DETECTOR="$2" + shift 2 + ;; + --full-image) + FULL_IMAGE_MODE=true + shift + ;; + --image-size) + if [ -z "$2" ]; then + echo "Error: --image-size requires a number argument" + exit 1 + fi + IMAGE_SIZE="$2" + shift 2 + ;; + --input) + if [ -z "$2" ]; then + echo "Error: --input requires a directory argument" + exit 1 + fi + INPUT_DIR="$2" + shift 2 + ;; + --target) + if [ -z "$2" ]; then + echo "Error: --target requires a directory argument" + exit 1 + fi + TARGET_DIR="$2" + shift 2 + ;; + --checkpoint-dir) + if [ -z "$2" ]; then + echo "Error: --checkpoint-dir requires a directory argument" + exit 1 + fi + CHECKPOINT_DIR="$2" + shift 2 + ;; + --validation-dir) + if [ -z "$2" ]; then + echo "Error: --validation-dir requires a directory argument" + exit 1 + fi + VALIDATION_DIR="$2" + shift 2 + ;; + --output-weights) + if [ -z "$2" ]; then + echo "Error: --output-weights requires a file path argument" + exit 1 + fi + OUTPUT_WEIGHTS="$2" + shift 2 + ;; + *) + echo "Unknown option: $1" + exit 1 + ;; + esac +done + +# Build training arguments +if [ "$FULL_IMAGE_MODE" = true ]; then + TRAINING_MODE_ARGS="--full-image --image-size $IMAGE_SIZE" +else + TRAINING_MODE_ARGS="--patch-size $PATCH_SIZE --patches-per-image $PATCHES_PER_IMAGE --detector $DETECTOR" +fi + +# Handle export-only mode +if [ "$EXPORT_ONLY" = true ]; then + echo "=== CNN v2 Export Weights Only ===" + echo "Checkpoint: $EXPORT_CHECKPOINT" + echo "" + + if [ ! -f "$EXPORT_CHECKPOINT" ]; then + echo "Error: Checkpoint file not found: $EXPORT_CHECKPOINT" + exit 1 + fi + + export_weights "$EXPORT_CHECKPOINT" "$OUTPUT_WEIGHTS" || { + echo "Error: Export failed" + exit 1 + } + + echo "" + echo "=== Export Complete ===" + echo "Output: $OUTPUT_WEIGHTS" + exit 0 +fi + +if [ "$VALIDATE_ONLY" = true ]; then + echo "=== CNN v2 Validation Only ===" + echo "Skipping training, using existing weights" + echo "" +else + echo "=== CNN v2 Complete Training Pipeline ===" + echo "Input: $INPUT_DIR" + echo "Target: $TARGET_DIR" + echo "Epochs: $EPOCHS" + echo "Checkpoint interval: $CHECKPOINT_EVERY" + echo "Mip level: $MIP_LEVEL (p0-p3 features)" + echo "" +fi + +if [ "$VALIDATE_ONLY" = false ]; then + # Step 1: Train model + echo "[1/4] Training CNN v2 model..." + +python3 "$SCRIPT_DIR/../training/train_cnn_v2.py" \ + --input "$INPUT_DIR" \ + --target "$TARGET_DIR" \ + $TRAINING_MODE_ARGS \ + --kernel-sizes "$KERNEL_SIZES" \ + --num-layers "$NUM_LAYERS" \ + --mip-level "$MIP_LEVEL" \ + --epochs "$EPOCHS" \ + --batch-size "$BATCH_SIZE" \ + --lr "$LEARNING_RATE" \ + --checkpoint-dir "$CHECKPOINT_DIR" \ + --checkpoint-every "$CHECKPOINT_EVERY" \ + $([ "$GRAYSCALE_LOSS" = true ] && echo "--grayscale-loss") + +if [ $? -ne 0 ]; then + echo "Error: Training failed" + exit 1 +fi + +echo "" +echo "Training complete!" +echo "" + +# Step 2: Export final checkpoint to shaders +FINAL_CHECKPOINT="$CHECKPOINT_DIR/checkpoint_epoch_${EPOCHS}.pth" + +if [ ! -f "$FINAL_CHECKPOINT" ]; then + echo "Warning: Final checkpoint not found, using latest available..." + FINAL_CHECKPOINT=$(find_latest_checkpoint) +fi + +if [ -z "$FINAL_CHECKPOINT" ] || [ ! -f "$FINAL_CHECKPOINT" ]; then + echo "Error: No checkpoint found in $CHECKPOINT_DIR" + exit 1 +fi + +echo "[2/4] Exporting final checkpoint to binary weights..." +echo "Checkpoint: $FINAL_CHECKPOINT" +export_weights "$FINAL_CHECKPOINT" "$OUTPUT_WEIGHTS" || { + echo "Error: Shader export failed" + exit 1 +} + +echo "" +fi # End of training/export section + +# Determine which checkpoint to use +if [ "$VALIDATE_ONLY" = true ]; then + FINAL_CHECKPOINT="${VALIDATE_CHECKPOINT:-$(find_latest_checkpoint)}" + echo "Using checkpoint: $FINAL_CHECKPOINT" + echo "" +fi + +# Step 3: Rebuild with new shaders +if [ "$VALIDATE_ONLY" = false ]; then + echo "[3/4] Rebuilding demo with new shaders..." + build_target demo64k || { + echo "Error: Build failed" + exit 1 + } + echo " → Build complete" + echo "" +fi + +# Step 4: Visual assessment - process final checkpoint only +if [ "$VALIDATE_ONLY" = true ]; then + echo "Validation on all input images (using existing weights)..." +else + echo "[4/4] Visual assessment on all input images..." +fi + +mkdir -p "$VALIDATION_DIR" +echo " Using checkpoint: $FINAL_CHECKPOINT" + +# Export weights for validation mode (already exported in step 2 for training mode) +if [ "$VALIDATE_ONLY" = true ]; then + export_weights "$FINAL_CHECKPOINT" "$OUTPUT_WEIGHTS" > /dev/null 2>&1 +fi + +# Build cnn_test +build_target cnn_test + +# Process all input images +echo -n " Processing images: " +for input_image in "$INPUT_DIR"/*.png; do + basename=$(basename "$input_image" .png) + echo -n "$basename " + build/cnn_test "$input_image" "$VALIDATION_DIR/${basename}_output.png" --weights "$OUTPUT_WEIGHTS" > /dev/null 2>&1 +done +echo "✓" + +# Build demo only if not in validate mode +[ "$VALIDATE_ONLY" = false ] && build_target demo64k + +echo "" +if [ "$VALIDATE_ONLY" = true ]; then + echo "=== Validation Complete ===" +else + echo "=== Training Pipeline Complete ===" +fi +echo "" +echo "Results:" +if [ "$VALIDATE_ONLY" = false ]; then + echo " - Checkpoints: $CHECKPOINT_DIR" + echo " - Final weights: $OUTPUT_WEIGHTS" +fi +echo " - Validation outputs: $VALIDATION_DIR" +echo "" +echo "Opening results directory..." +open "$VALIDATION_DIR" 2>/dev/null || xdg-open "$VALIDATION_DIR" 2>/dev/null || true + +if [ "$VALIDATE_ONLY" = false ]; then + echo "" + echo "Run demo to see final result:" + echo " ./build/demo64k" +fi diff --git a/cnn_v2/shaders/cnn_v2_compute.wgsl b/cnn_v2/shaders/cnn_v2_compute.wgsl new file mode 100644 index 0000000..cdbfd74 --- /dev/null +++ b/cnn_v2/shaders/cnn_v2_compute.wgsl @@ -0,0 +1,143 @@ +// CNN v2 Compute Shader - Uniform 12D→4D Architecture +// All layers: input/previous (4D) + static (8D) = 12D → 4 channels +// Storage buffer weights, ping-pong execution +// Per-layer kernel sizes supported via LayerParams + +// Push constants for layer parameters (passed per dispatch) +struct LayerParams { + kernel_size: u32, + in_channels: u32, + out_channels: u32, + weight_offset: u32, // Offset in f16 units + is_output_layer: u32, // 1 if final layer (sigmoid), 0 otherwise (relu) + blend_amount: f32, // [0,1] blend with original + is_layer_0: u32, // 1 if first layer (clamp [0,1]), 0 otherwise +} + +@group(0) @binding(0) var static_features: texture_2d; // 8D static features (p0-p3 + spatial) +@group(0) @binding(1) var layer_input: texture_2d; // 4D previous/input (RGBD or prev layer) +@group(0) @binding(2) var output_tex: texture_storage_2d; // 4D output +@group(0) @binding(3) var weights_buffer: array; // Packed f16 weights +@group(0) @binding(4) var params: LayerParams; +@group(0) @binding(5) var original_input: texture_2d; // Original RGB for blending + +fn unpack_static_features(coord: vec2) -> array { + let packed = textureLoad(static_features, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn unpack_layer_channels(coord: vec2) -> vec4 { + let packed = textureLoad(layer_input, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + return vec4(v0.x, v0.y, v1.x, v1.y); +} + +fn pack_channels(values: vec4) -> vec4 { + return vec4( + pack2x16float(vec2(values.x, values.y)), + pack2x16float(vec2(values.z, values.w)), + 0u, // Unused + 0u // Unused + ); +} + +// Get weight from storage buffer (f16 packed as u32 pairs) +// Buffer layout: [header: 4 u32][layer_info: N×5 u32][weights: packed f16] +// TODO: Support 8-bit quantized weights (4× per u32) for 2× size reduction +fn get_weight(idx: u32) -> f32 { + // Skip header (16 bytes = 4 u32) and layer info + // Weights start after header + layer_info, but weight_offset already accounts for this + let pair_idx = idx / 2u; + let packed = weights_buffer[pair_idx]; + let unpacked = unpack2x16float(packed); + return select(unpacked.y, unpacked.x, (idx & 1u) == 0u); +} + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) { + let coord = vec2(id.xy); + let dims = textureDimensions(static_features); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) { + return; + } + + let kernel_size = params.kernel_size; + let in_channels = params.in_channels; // Always 12 (4 prev + 8 static) + let out_channels = params.out_channels; // Always 4 + let weight_offset = params.weight_offset; + let is_output = params.is_output_layer != 0u; + + let kernel_radius = i32(kernel_size / 2u); + + // Load static features (8D) and previous/input layer (4D) + let static_feat = unpack_static_features(coord); + + // Convolution: 12D input → 4D output + var output: vec4 = vec4(0.0); + for (var c: u32 = 0u; c < 4u; c++) { + var sum: f32 = 0.0; + + // Convolve over kernel + for (var ky: i32 = -kernel_radius; ky <= kernel_radius; ky++) { + for (var kx: i32 = -kernel_radius; kx <= kernel_radius; kx++) { + let sample_coord = coord + vec2(kx, ky); + + // Border handling (clamp) + let clamped = vec2( + clamp(sample_coord.x, 0, i32(dims.x) - 1), + clamp(sample_coord.y, 0, i32(dims.y) - 1) + ); + + // Load features at this spatial location + let static_local = unpack_static_features(clamped); + let layer_local = unpack_layer_channels(clamped); // 4D + + // Weight index calculation + let ky_idx = u32(ky + kernel_radius); + let kx_idx = u32(kx + kernel_radius); + let spatial_idx = ky_idx * kernel_size + kx_idx; + + // Accumulate: previous/input channels (4D) + for (var i: u32 = 0u; i < 4u; i++) { + let w_idx = weight_offset + + c * 12u * kernel_size * kernel_size + + i * kernel_size * kernel_size + spatial_idx; + sum += get_weight(w_idx) * layer_local[i]; + } + + // Accumulate: static features (8D) + for (var i: u32 = 0u; i < 8u; i++) { + let w_idx = weight_offset + + c * 12u * kernel_size * kernel_size + + (4u + i) * kernel_size * kernel_size + spatial_idx; + sum += get_weight(w_idx) * static_local[i]; + } + } + } + + // Activation (matches train_cnn_v2.py) + if (is_output || params.is_layer_0 != 0u) { + output[c] = 1.0 / (1.0 + exp(-sum)); // Sigmoid [0,1] + } else { + output[c] = max(0.0, sum); // ReLU + } + } + + // Blend with original on final layer + if (is_output) { + let original = textureLoad(original_input, coord, 0).rgb; + let result_rgb = vec3(output.x, output.y, output.z); + let blended = mix(original, result_rgb, params.blend_amount); + output.x = blended.r; + output.y = blended.g; + output.z = blended.b; + } + + textureStore(output_tex, coord, pack_channels(output)); +} diff --git a/cnn_v2/shaders/cnn_v2_layer_0.wgsl b/cnn_v2/shaders/cnn_v2_layer_0.wgsl new file mode 100644 index 0000000..8e14957 --- /dev/null +++ b/cnn_v2/shaders/cnn_v2_layer_0.wgsl @@ -0,0 +1,174 @@ +// CNN v2 Layer 0 - Auto-generated +// Kernel: 3×3, In: 8, Out: 8 + +const KERNEL_SIZE: u32 = 3u; +const IN_CHANNELS: u32 = 8u; +const OUT_CHANNELS: u32 = 8u; +const KERNEL_RADIUS: i32 = 1; + +// Weights quantized to float16 (stored as f32 in WGSL) +const weights: array = array( + 0.057281, -0.041962, 0.003933, 0.026459, 0.304199, 0.067261, 0.191895, 0.047455, + 0.074402, 0.201660, 0.158325, 0.150513, 0.219238, 0.260010, 0.319336, 0.208618, + 0.050201, 0.090210, 0.086853, 0.181152, 0.060486, 0.167847, 0.161499, 0.265869, + 0.163818, 0.100647, 0.243408, -0.008553, -0.010849, 0.046509, -0.060608, -0.022263, + 0.094360, -0.043854, -0.005329, -0.093262, 0.032349, 0.007259, 0.039948, -0.018692, + -0.000618, 0.052368, -0.038055, 0.118042, -0.084595, 0.044281, -0.107056, 0.089478, + -0.076477, 0.017441, 0.088135, 0.076721, -0.063965, 0.001612, 0.062469, 0.067505, + 0.035736, 0.115051, -0.117737, -0.076843, -0.008888, -0.002028, -0.061005, 0.081726, + 0.115051, -0.028183, 0.043213, -0.079285, -0.040314, -0.047699, -0.051575, -0.052521, + 0.071533, 0.084656, 0.051910, 0.090637, -0.104248, -0.066467, -0.032104, -0.006977, + 0.075439, -0.004841, 0.084656, -0.034698, 0.035675, -0.101929, -0.035034, -0.036804, + 0.069641, -0.026840, -0.017807, -0.088318, -0.125000, -0.042847, -0.003063, 0.007622, + 0.076416, 0.094971, -0.019058, 0.083496, -0.085205, 0.036285, -0.077209, 0.082458, + 0.056549, 0.038818, 0.092224, -0.002499, 0.069641, 0.097229, 0.069275, -0.111084, + -0.092041, -0.020462, -0.061279, -0.032196, -0.088623, 0.032227, -0.117004, -0.125854, + -0.015884, 0.093018, -0.070923, -0.117615, -0.081848, -0.115479, 0.033508, -0.026443, + -0.009850, -0.063232, 0.098328, -0.000984, 0.039886, -0.085754, -0.108826, 0.030258, + 0.091675, 0.024384, -0.118958, -0.077148, -0.122437, -0.002090, -0.089539, 0.096741, + 0.095337, 0.108582, -0.101807, 0.152222, 0.206177, 0.050323, -0.111450, -0.104431, + -0.037445, 0.276611, 0.244019, 0.171143, 0.131592, 0.056030, 0.141602, 0.014267, + -0.025955, -0.019730, 0.155884, 0.072144, 0.176636, -0.010117, 0.141724, 0.103027, + -0.253174, -0.229370, -0.105713, -0.005898, 0.075439, -0.002014, -0.010506, -0.108093, + -0.016724, 0.108215, 0.053589, -0.044586, 0.030396, -0.077759, 0.058594, -0.018463, + 0.027100, 0.030823, -0.026947, -0.014084, 0.121643, 0.116638, -0.010239, 0.106262, + -0.109070, -0.044281, -0.045319, -0.021942, 0.083923, 0.114929, 0.154541, 0.078186, + -0.047394, 0.007957, 0.099182, -0.030075, 0.103699, 0.080994, -0.085144, 0.047180, + 0.099792, 0.081116, 0.084961, 0.151123, 0.000963, 0.029221, 0.073181, 0.086609, + 0.149048, -0.052185, -0.158936, 0.146240, 0.020004, 0.063110, 0.111877, 0.037201, + 0.087585, 0.134277, 0.058258, -0.075256, 0.141357, 0.045776, 0.171753, 0.186035, + 0.093201, 0.202637, 0.018723, -0.047638, 0.072510, 0.132812, 0.182251, 0.191650, + 0.163818, 0.146362, 0.124451, -0.082214, 0.094482, -0.007275, 0.029099, -0.040314, + -0.017624, -0.018860, -0.108398, -0.111145, 0.058289, -0.106995, -0.091919, 0.069824, + -0.084045, -0.105957, 0.065002, -0.012894, 0.042297, -0.081299, -0.112976, 0.012314, + 0.015625, -0.100708, -0.039673, 0.092041, 0.037201, 0.089722, 0.064087, 0.000403, + 0.120667, -0.012238, -0.055695, 0.010620, -0.022110, -0.008751, 0.038605, 0.075256, + 0.041260, 0.128296, -0.072021, 0.020828, -0.072449, 0.051239, 0.034058, 0.122803, + -0.062103, 0.156006, -0.111633, 0.043671, 0.209229, 0.006088, 0.141968, 0.209961, + 0.122620, -0.004547, 0.107727, 0.115601, 0.003378, 0.375732, 0.068481, 0.037842, + 0.159546, -0.014450, 0.073425, 0.168701, -0.052643, 0.060699, 0.333740, 0.033905, + -0.060150, 0.053558, 0.165527, -0.052460, -0.047882, 0.080750, 0.110352, -0.057098, + 0.057983, -0.018692, 0.019714, -0.056427, -0.053314, -0.001763, 0.027039, 0.003395, + -0.131226, -0.068481, -0.086609, 0.065186, 0.084717, 0.036530, 0.043488, 0.013893, + -0.076660, 0.081177, 0.037476, -0.124084, -0.070312, -0.027130, -0.009331, -0.128174, + -0.075256, 0.098206, -0.046539, -0.045319, 0.083923, -0.050598, 0.063477, 0.007408, + 0.026794, -0.090454, -0.083435, 0.129761, 0.044556, 0.051849, 0.115662, 0.071167, + 0.004414, 0.048035, -0.148682, 0.098938, 0.200562, 0.111938, 0.208496, 0.200684, + -0.050262, 0.119568, 0.062988, 0.072083, 0.123779, 0.369629, 0.317627, 0.187622, + 0.157227, 0.183960, 0.031921, 0.142944, 0.080627, 0.218628, 0.264160, 0.156128, + 0.084961, 0.029343, 0.057617, 0.089233, 0.041138, 0.044373, 0.074707, 0.025818, + 0.113708, -0.045380, -0.114929, 0.104370, -0.012238, -0.174194, -0.169312, -0.070312, + -0.005863, 0.027481, 0.053345, -0.016006, -0.057953, -0.010284, 0.034241, -0.041077, + -0.002373, 0.034515, 0.078552, -0.066162, -0.035400, 0.072510, 0.060425, -0.037720, + -0.025955, 0.118042, -0.071777, 0.133667, 0.012192, -0.080933, 0.093445, 0.052826, + -0.037354, -0.052277, 0.124084, 0.029861, 0.137085, 0.053009, -0.034180, -0.011421, + 0.089233, 0.172729, 0.146118, 0.003944, 0.279541, 0.162842, 0.112244, 0.204956, + 0.059753, 0.117737, 0.330322, 0.185547, 0.194946, 0.404541, 0.274658, 0.177612, + 0.153320, 0.189575, 0.032257, 0.285400, 0.158203, 0.048035, 0.476562, 0.301025, + -0.179565, 0.160767, 0.137207, 0.102478, -0.060547, 0.060364, -0.091858, 0.064209, + 0.082642, 0.044769, -0.096436, -0.103699, -0.021683, 0.007221, -0.048737, 0.071228, + -0.069580, 0.066528, -0.122864, -0.008415, -0.094788, 0.040131, -0.091431, -0.029602, + -0.112488, -0.074158, -0.004898, -0.006721, -0.118286, -0.047516, 0.069519, 0.121521, + -0.004158, 0.167603, -0.092468, -0.049927, 0.006599, 0.097595, 0.064087, 0.083435, + 0.026993, 0.071411, 0.020538, 0.022293, 0.022858, 0.124268, 0.098999, -0.031738, + 0.019806, -0.087341, -0.096558, -0.099304, -0.113159, 0.021744, -0.080200, -0.056030, + 0.089661, -0.055115, -0.115845, -0.040222, 0.035919, 0.027832, 0.034668, 0.072632, + 0.071838, -0.081116, 0.050262, -0.037872, 0.054047, -0.096680, -0.102051, -0.044281, + 0.078796, -0.095154, -0.013229, 0.031555, -0.058533, -0.114441, -0.008530, 0.112732, + -0.057251, 0.096191, -0.008385, 0.052246, -0.016983, 0.092041, 0.013710, 0.012299, + -0.109497, 0.025604, -0.121643, -0.023819, 0.039490, -0.090088, -0.013145, -0.101562, + -0.115051, 0.050232, -0.047119, -0.055847, -0.017563, 0.103760, 0.116333, -0.061768, + -0.083069, -0.030319, 0.078003, -0.010124, 0.044617, -0.045868, 0.103638, 0.032379, + -0.093506, -0.048004, -0.022079, -0.004353, -0.048187, -0.025330, -0.070740, -0.014671 +); + +@group(0) @binding(0) var static_features: texture_2d; +@group(0) @binding(1) var layer_input: texture_2d; +@group(0) @binding(2) var output_tex: texture_storage_2d; + +fn unpack_static_features(coord: vec2) -> array { + let packed = textureLoad(static_features, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn unpack_layer_channels(coord: vec2) -> array { + let packed = textureLoad(layer_input, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn pack_channels(values: array) -> vec4 { + return vec4( + pack2x16float(vec2(values[0], values[1])), + pack2x16float(vec2(values[2], values[3])), + pack2x16float(vec2(values[4], values[5])), + pack2x16float(vec2(values[6], values[7])) + ); +} + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) { + let coord = vec2(id.xy); + let dims = textureDimensions(static_features); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) { + return; + } + + // Load static features (always available) + let static_feat = unpack_static_features(coord); + + // Convolution + var output: array; + for (var c: u32 = 0u; c < OUT_CHANNELS; c++) { + var sum: f32 = 0.0; + + for (var ky: i32 = -KERNEL_RADIUS; ky <= KERNEL_RADIUS; ky++) { + for (var kx: i32 = -KERNEL_RADIUS; kx <= KERNEL_RADIUS; kx++) { + let sample_coord = coord + vec2(kx, ky); + + // Border handling (clamp) + let clamped = vec2( + clamp(sample_coord.x, 0, i32(dims.x) - 1), + clamp(sample_coord.y, 0, i32(dims.y) - 1) + ); + + // Load input features + let static_local = unpack_static_features(clamped); + let layer_local = unpack_layer_channels(clamped); + + // Weight index calculation + let ky_idx = u32(ky + KERNEL_RADIUS); + let kx_idx = u32(kx + KERNEL_RADIUS); + let spatial_idx = ky_idx * KERNEL_SIZE + kx_idx; + + // Accumulate: static features (8D) + for (var i: u32 = 0u; i < 8u; i++) { + let w_idx = c * IN_CHANNELS * KERNEL_SIZE * KERNEL_SIZE + + i * KERNEL_SIZE * KERNEL_SIZE + spatial_idx; + sum += weights[w_idx] * static_local[i]; + } + + // Accumulate: layer input channels (if layer_idx > 0) + let prev_channels = IN_CHANNELS - 8u; + for (var i: u32 = 0u; i < prev_channels; i++) { + let w_idx = c * IN_CHANNELS * KERNEL_SIZE * KERNEL_SIZE + + (8u + i) * KERNEL_SIZE * KERNEL_SIZE + spatial_idx; + sum += weights[w_idx] * layer_local[i]; + } + } + } + + output[c] = max(0.0, sum); // ReLU + } + + // Pack and store + textureStore(output_tex, coord, pack_channels(output)); +} diff --git a/cnn_v2/shaders/cnn_v2_layer_1.wgsl b/cnn_v2/shaders/cnn_v2_layer_1.wgsl new file mode 100644 index 0000000..f490d13 --- /dev/null +++ b/cnn_v2/shaders/cnn_v2_layer_1.wgsl @@ -0,0 +1,174 @@ +// CNN v2 Layer 1 - Auto-generated +// Kernel: 3×3, In: 16, Out: 4 + +const KERNEL_SIZE: u32 = 3u; +const IN_CHANNELS: u32 = 16u; +const OUT_CHANNELS: u32 = 4u; +const KERNEL_RADIUS: i32 = 1; + +// Weights quantized to float16 (stored as f32 in WGSL) +const weights: array = array( + 0.337402, 0.638672, -0.481201, 0.699707, 1.127930, -0.018280, -0.062195, 0.148682, + -0.655273, 0.448975, 0.969238, -0.280762, 0.817383, 1.271484, 0.421387, -0.163696, + 0.305664, -0.454834, 0.354004, 0.932617, -0.411377, 0.581543, 1.263672, 0.422363, + -0.380371, 0.152588, -0.668945, -0.063782, 0.060730, 0.022018, -0.075195, -0.049286, + 0.068542, 0.057343, -0.009773, 0.006344, -0.080872, -0.179932, -0.297119, 0.098328, + 0.061951, -0.088989, 0.047913, 0.093628, -0.091858, -0.068298, 0.102600, -0.044067, + -0.054230, -0.031799, 0.050934, -0.300049, -0.202637, -0.203613, -0.294189, -0.361084, + 0.277344, -0.213257, -0.239624, 0.193237, -0.215210, -0.295166, 0.298828, -0.065369, + 0.148926, 0.024963, 0.272705, 0.368164, 0.173096, 0.061279, 0.291260, 0.151611, + 0.411133, 0.216431, -0.179932, 0.506348, 0.319580, 0.059875, -0.134399, -0.150635, + -0.275391, 0.029480, 0.115417, 0.063782, 0.018723, -0.073364, -0.019653, 0.066467, + -0.086731, 0.113220, 0.110535, 0.011940, -0.094727, 0.262207, 0.180298, 0.141357, + 0.249634, 0.199585, 0.120605, 0.403809, 0.242676, -0.028442, 0.251953, 0.130737, + 0.152832, -0.306396, -0.324951, -0.176514, 0.161133, 0.333252, -0.195068, 0.250244, + 0.569824, 0.011223, -0.186035, 0.048279, -0.325439, 0.272217, 0.144043, -0.142700, + 0.447754, 0.434082, 0.124878, -0.157471, -0.120422, -0.281494, 0.338135, 0.266113, + -0.301514, 0.424805, 0.541504, -0.195679, 0.054962, 0.061798, -0.323975, 0.056732, + 0.072571, -0.087341, 0.052856, -0.057220, 0.023270, 0.071472, 0.014038, 0.083008, + -0.050659, 0.020111, 0.035614, -0.038086, -0.042786, 0.060242, -0.050079, -0.044403, + -0.059631, 0.075500, 0.056000, 0.010910, -0.064026, -0.016037, -0.050720, 0.050171, + -0.075256, -0.014183, 0.047058, -0.086731, 0.027939, 0.063232, -0.024597, -0.039551, + 0.000622, -0.048370, -0.001906, 0.058868, -0.074524, 0.019714, -0.036011, 0.028442, + 0.009766, -0.060577, -0.007416, -0.014381, 0.002317, -0.023483, 0.014313, 0.057434, + 0.063110, 0.030350, -0.027557, 0.023270, 0.055115, -0.003502, 0.012268, -0.054993, + -0.084961, -0.022736, 0.076233, 0.027573, -0.068787, -0.036987, -0.018539, -0.049347, + 0.032227, 0.033081, 0.050476, 0.043030, 0.023636, -0.039764, -0.018600, 0.073669, + 0.032166, -0.047119, -0.033325, -0.038605, 0.034119, -0.076843, 0.005863, -0.049103, + 0.065796, -0.056458, 0.054504, -0.008354, -0.018509, -0.057739, -0.075684, -0.053680, + 0.036804, 0.020721, -0.056183, 0.021774, -0.043884, 0.033661, -0.029633, 0.027374, + -0.087891, 0.030853, -0.040070, 0.013733, -0.082275, -0.072571, -0.055756, 0.002262, + 0.004421, -0.012169, -0.078064, -0.063904, -0.051758, -0.033264, -0.059265, -0.062256, + 0.063782, -0.088745, -0.026855, 0.062805, -0.036591, 0.037659, -0.012970, 0.025513, + -0.000908, 0.027084, 0.001842, -0.080750, -0.049713, -0.069397, -0.046448, -0.031006, + 0.012543, 0.009369, -0.080139, -0.034363, 0.003361, -0.052704, 0.041870, 0.059265, + 0.029938, 0.000138, 0.049896, 0.068787, 0.040405, -0.073608, 0.047668, 0.015320, + -0.033203, -0.016983, 0.034149, -0.010323, 0.029877, 0.078003, -0.054688, -0.021805, + -0.019409, 0.010284, 0.089172, -0.050385, 0.024857, -0.041992, 0.016602, 0.082397, + 0.081970, 0.096375, 0.060760, -0.006603, 0.029907, 0.012131, 0.104980, 0.034210, + 0.074707, -0.028320, -0.020248, 0.114868, -0.036957, 0.040192, 0.002888, 0.034973, + -0.038635, -0.018204, -0.058563, 0.029419, 0.013344, 0.027618, 0.073669, -0.038361, + 0.080933, 0.044586, -0.013214, 0.022675, 0.084351, 0.081848, 0.027328, 0.043915, + 0.040771, 0.078918, 0.054443, -0.049652, 0.073547, 0.103882, 0.065918, 0.070923, + -0.037476, -0.011215, -0.021408, 0.094727, 0.042450, 0.032806, -0.064026, 0.023941, + 0.011780, 0.041260, -0.038818, 0.079163, 0.079468, 0.053680, 0.047150, 0.003571, + 0.054840, 0.045929, -0.041382, -0.033539, 0.069153, 0.046234, 0.119263, -0.006340, + -0.050323, 0.030212, 0.069092, 0.045441, 0.096313, -0.024628, -0.088745, 0.009033, + -0.016830, 0.028534, -0.042755, -0.031921, 0.013611, -0.029251, -0.051483, -0.005848, + -0.032837, -0.058136, 0.075989, -0.008125, 0.108765, -0.004745, -0.003422, 0.079590, + 0.090515, -0.019196, -0.006786, 0.059479, -0.041168, 0.093445, 0.075439, -0.025055, + 0.067139, 0.011734, 0.031586, 0.029587, 0.098267, 0.025848, 0.095276, 0.003189, + 0.105408, 0.018799, -0.102478, 0.033813, 0.004272, 0.020477, 0.033142, 0.009727, + -0.021393, 0.120300, 0.088684, -0.037842, -0.094177, 0.017944, 0.020126, -0.002304, + -0.016006, 0.018112, 0.072693, -0.072021, -0.171265, -0.053528, -0.093201, 0.024124, + -0.050476, -0.023422, -0.071167, 0.046478, 0.034607, 0.076904, 0.013077, -0.082031, + 0.091858, -0.001575, 0.083801, 0.078003, 0.019119, -0.004967, 0.027298, 0.027740, + 0.032623, 0.048370, 0.029099, 0.093201, 0.049957, -0.007191, 0.059631, 0.008659, + 0.042725, -0.009369, 0.089417, 0.074951, -0.024704, 0.005344, 0.123840, 0.080322, + 0.096375, 0.070312, -0.010399, 0.033203, -0.009743, -0.030045, -0.039520, 0.042023, + -0.017441, 0.073486, 0.049500, -0.039734, 0.009811, 0.093262, -0.069641, 0.099365, + -0.010414, 0.048859, 0.099182, -0.007256, -0.023941, -0.021393, -0.005703, 0.025055, + 0.054535, 0.093384, -0.033661, 0.073242, 0.055023, 0.037170, -0.009300, 0.048615, + 0.019150, 0.019409, -0.080688, -0.050049, 0.104126, -0.023193, 0.044708, 0.111816, + 0.061584, 0.042755, -0.013863, -0.008385, -0.039703, 0.070618, -0.016922, -0.040833, + 0.051178, -0.060333, -0.004368, -0.009827, 0.051544, 0.072083, 0.068176, 0.148071, + 0.159424, 0.017578, 0.089905, -0.006794, 0.066101, -0.051117, 0.088684, -0.002989, + -0.066895, 0.089844, 0.012131, -0.020203, 0.011230, 0.000327, 0.073669, 0.060669, + 0.091064, 0.075989, 0.051971, 0.045044, 0.033875, 0.040466, -0.029449, 0.128418, + -0.000229, -0.026901, 0.052063, 0.000995, -0.032532, 0.105896, -0.001241, 0.114075, + 0.047607, 0.090332, 0.063660, 0.016495, 0.124817, 0.090942, 0.021545, 0.007164, + 0.074890, 0.118347, 0.047394, 0.052856, 0.104980, 0.009384, 0.034363, 0.019073, + 0.072388, -0.013313, 0.119141, 0.021255, 0.103210, 0.058319, 0.186035, -0.010818, + 0.037109, -0.044037, -0.075989, -0.001281, 0.017899, 0.030701, -0.080261, 0.082703 +); + +@group(0) @binding(0) var static_features: texture_2d; +@group(0) @binding(1) var layer_input: texture_2d; +@group(0) @binding(2) var output_tex: texture_storage_2d; + +fn unpack_static_features(coord: vec2) -> array { + let packed = textureLoad(static_features, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn unpack_layer_channels(coord: vec2) -> array { + let packed = textureLoad(layer_input, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn pack_channels(values: array) -> vec4 { + return vec4( + pack2x16float(vec2(values[0], values[1])), + pack2x16float(vec2(values[2], values[3])), + pack2x16float(vec2(values[4], values[5])), + pack2x16float(vec2(values[6], values[7])) + ); +} + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) { + let coord = vec2(id.xy); + let dims = textureDimensions(static_features); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) { + return; + } + + // Load static features (always available) + let static_feat = unpack_static_features(coord); + + // Convolution + var output: array; + for (var c: u32 = 0u; c < OUT_CHANNELS; c++) { + var sum: f32 = 0.0; + + for (var ky: i32 = -KERNEL_RADIUS; ky <= KERNEL_RADIUS; ky++) { + for (var kx: i32 = -KERNEL_RADIUS; kx <= KERNEL_RADIUS; kx++) { + let sample_coord = coord + vec2(kx, ky); + + // Border handling (clamp) + let clamped = vec2( + clamp(sample_coord.x, 0, i32(dims.x) - 1), + clamp(sample_coord.y, 0, i32(dims.y) - 1) + ); + + // Load input features + let static_local = unpack_static_features(clamped); + let layer_local = unpack_layer_channels(clamped); + + // Weight index calculation + let ky_idx = u32(ky + KERNEL_RADIUS); + let kx_idx = u32(kx + KERNEL_RADIUS); + let spatial_idx = ky_idx * KERNEL_SIZE + kx_idx; + + // Accumulate: static features (8D) + for (var i: u32 = 0u; i < 8u; i++) { + let w_idx = c * IN_CHANNELS * KERNEL_SIZE * KERNEL_SIZE + + i * KERNEL_SIZE * KERNEL_SIZE + spatial_idx; + sum += weights[w_idx] * static_local[i]; + } + + // Accumulate: layer input channels (if layer_idx > 0) + let prev_channels = IN_CHANNELS - 8u; + for (var i: u32 = 0u; i < prev_channels; i++) { + let w_idx = c * IN_CHANNELS * KERNEL_SIZE * KERNEL_SIZE + + (8u + i) * KERNEL_SIZE * KERNEL_SIZE + spatial_idx; + sum += weights[w_idx] * layer_local[i]; + } + } + } + + output[c] = max(0.0, sum); // ReLU + } + + // Pack and store + textureStore(output_tex, coord, pack_channels(output)); +} diff --git a/cnn_v2/shaders/cnn_v2_layer_2.wgsl b/cnn_v2/shaders/cnn_v2_layer_2.wgsl new file mode 100644 index 0000000..2f9836a --- /dev/null +++ b/cnn_v2/shaders/cnn_v2_layer_2.wgsl @@ -0,0 +1,156 @@ +// CNN v2 Layer 2 - Auto-generated +// Kernel: 3×3, In: 12, Out: 4 + +const KERNEL_SIZE: u32 = 3u; +const IN_CHANNELS: u32 = 12u; +const OUT_CHANNELS: u32 = 4u; +const KERNEL_RADIUS: i32 = 1; + +// Weights quantized to float16 (stored as f32 in WGSL) +const weights: array = array( + 0.030212, -0.041351, 0.053864, -0.025635, 0.099976, -0.016830, -0.068665, 0.112488, + -0.069824, 0.030197, 0.020142, 0.101807, 0.061920, 0.022415, -0.025864, -0.056366, + 0.085571, -0.053650, 0.109802, 0.129272, 0.023438, 0.087341, 0.066284, 0.037079, + -0.067566, 0.021530, -0.046814, 0.029343, -0.028534, 0.047150, -0.079346, -0.022675, + -0.019669, -0.024185, 0.029587, 0.068970, 0.108826, 0.050598, -0.072144, 0.083008, + -0.002201, 0.006275, 0.056396, 0.001884, 0.097168, -0.028503, -0.002499, 0.008919, + -0.013771, -0.017502, -0.033478, 0.105530, 0.032898, 0.068726, -0.036285, -0.021011, + -0.018250, 0.073914, 0.024277, 0.061066, 0.008682, -0.022766, 0.074219, 0.094421, + 0.050903, 0.072571, 0.117493, -0.033234, 0.067993, -0.008049, 0.046997, -0.064209, + -0.381104, 0.107788, -0.213867, 0.145142, 0.514160, 0.407715, -0.317871, 0.249023, + 0.055634, -0.006294, -0.067444, 0.025131, 0.012939, -0.074158, -0.013741, -0.033020, + 0.026871, -0.007671, 0.089661, -0.003016, 0.029007, -0.038483, 0.045044, 0.104065, + 0.077148, 0.092468, -0.090027, -0.048126, 0.096863, -0.088013, 0.009483, 0.075012, + -0.076843, -0.085449, -0.066040, 0.019165, -0.019958, 0.083496, 0.069275, -0.019714, + 0.027786, -0.042389, 0.054718, 0.010635, -0.071777, 0.029282, -0.003605, 0.113770, + 0.080994, 0.106079, 0.047333, -0.013733, 0.034760, 0.099365, -0.020813, 0.095886, + 0.052490, -0.049194, 0.047394, 0.072510, -0.030930, -0.003782, -0.038025, -0.019318, + -0.047852, -0.043915, 0.026810, -0.041138, 0.038422, 0.009605, -0.080688, -0.019653, + 0.075256, -0.013817, -0.022400, 0.050629, 0.048462, 0.072998, -0.009109, 0.070923, + 0.079895, 0.071350, 0.002869, 0.081543, 0.037231, 0.020767, -0.017929, 0.042328, + -0.075134, -0.010681, -0.009079, 0.057007, -0.040253, -0.025574, -0.041534, 0.105835, + -0.039703, 0.032104, 0.076050, 0.070923, -0.013046, -0.054108, -0.024582, -0.033997, + 0.092285, 0.000525, 0.114685, 0.036926, -0.419434, 0.087891, -0.187866, 0.128906, + 0.665527, 0.268311, -0.337891, 0.195557, 0.140503, 0.014465, -0.043671, 0.031677, + 0.073059, 0.085144, 0.014290, -0.046967, 0.033356, 0.004177, 0.102844, 0.015259, + 0.026627, -0.005032, 0.111694, -0.010590, 0.029816, 0.108154, -0.072327, 0.056213, + 0.022903, 0.053772, 0.084473, -0.059845, -0.032776, -0.000015, -0.093872, -0.085815, + 0.081604, 0.069336, 0.034149, -0.067322, -0.020859, 0.120911, 0.077209, -0.016388, + 0.050140, -0.045563, -0.046326, 0.032623, -0.005009, 0.008003, 0.109192, 0.086548, + 0.096558, 0.118530, 0.035034, 0.110352, -0.041748, 0.009178, 0.049957, 0.084839, + 0.042053, -0.069153, -0.024796, -0.094604, -0.047028, -0.053802, 0.024979, 0.049591, + -0.016373, -0.047607, -0.008797, -0.058868, 0.107178, 0.055695, 0.092407, 0.092346, + 0.053894, 0.054657, -0.039703, -0.073792, 0.041779, -0.044159, 0.099182, 0.037109, + 0.097778, 0.098206, -0.057831, -0.054016, -0.068604, -0.061584, -0.054382, 0.005268, + 0.096008, -0.007118, -0.063049, 0.059113, 0.076904, 0.045288, -0.055695, -0.052612, + -0.022110, 0.049103, 0.095276, 0.014572, 0.064819, 0.014671, 0.029800, 0.066284, + -0.383301, 0.071838, -0.207275, 0.099365, 0.640137, 0.393311, -0.334229, 0.275391, + -0.013977, -0.025269, -0.007065, -0.033478, -0.017349, 0.026764, 0.005192, 0.093384, + 0.014313, 0.018906, 0.006962, 0.094849, 0.005390, 0.101624, -0.041199, 0.026245, + 0.027588, 0.062408, 0.033356, -0.010826, 0.067993, -0.054199, 0.076416, 0.023315, + -0.002886, -0.112061, -0.041473, -0.012703, 0.016022, 0.010506, -0.021362, -0.037750, + 0.062927, 0.061920, 0.038177, -0.037201, -0.011620, 0.014015, -0.062164, -0.045441, + -0.063416, -0.040100, 0.035950, 0.045563, -0.017227, -0.060547, -0.017593, 0.111877, + 0.121521, 0.073853, 0.023331, -0.012428, 0.018478, -0.010948, 0.030716, 0.043427, + 0.003117, -0.069092, 0.038361, -0.053497, 0.039154, -0.085754, 0.012642, -0.051208, + 0.022934, 0.127197, 0.117920, 0.074036, 0.083313, -0.061951, 0.079224, 0.091248, + 0.009132, 0.069946, 0.123474, 0.130127, 0.118835, 0.020874, -0.045380, -0.000111, + 0.111206, 0.054688, 0.008995, 0.085693, 0.005562, 0.103088, -0.034698, 0.119934, + -0.067200, 0.065430, -0.021942, 0.089783, 0.033112, -0.025467, 0.040161, -0.052155, + -0.048920, 0.031250, 0.112549, 0.122192, 0.126587, 0.180908, 0.194946, 0.121704, + 0.217529, 0.224243, 0.269287, 0.222656, 0.288086, 0.035492, 0.066711, -0.046600, + 0.085144, 0.013855, -0.065979, -0.083252, -0.058289, 0.104126, 0.013702, -0.018188, + 0.036591, 0.099854, 0.056061, 0.151855, 0.062134, 0.133789, 0.084045, 0.095825, + 0.036987, 0.022308, 0.070923, 0.031036, 0.101868, 0.062347, 0.141235, 0.066650 +); + +@group(0) @binding(0) var static_features: texture_2d; +@group(0) @binding(1) var layer_input: texture_2d; +@group(0) @binding(2) var output_tex: texture_storage_2d; + +fn unpack_static_features(coord: vec2) -> array { + let packed = textureLoad(static_features, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn unpack_layer_channels(coord: vec2) -> array { + let packed = textureLoad(layer_input, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn pack_channels(values: array) -> vec4 { + return vec4( + pack2x16float(vec2(values[0], values[1])), + pack2x16float(vec2(values[2], values[3])), + pack2x16float(vec2(values[4], values[5])), + pack2x16float(vec2(values[6], values[7])) + ); +} + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) { + let coord = vec2(id.xy); + let dims = textureDimensions(static_features); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) { + return; + } + + // Load static features (always available) + let static_feat = unpack_static_features(coord); + + // Convolution + var output: array; + for (var c: u32 = 0u; c < OUT_CHANNELS; c++) { + var sum: f32 = 0.0; + + for (var ky: i32 = -KERNEL_RADIUS; ky <= KERNEL_RADIUS; ky++) { + for (var kx: i32 = -KERNEL_RADIUS; kx <= KERNEL_RADIUS; kx++) { + let sample_coord = coord + vec2(kx, ky); + + // Border handling (clamp) + let clamped = vec2( + clamp(sample_coord.x, 0, i32(dims.x) - 1), + clamp(sample_coord.y, 0, i32(dims.y) - 1) + ); + + // Load input features + let static_local = unpack_static_features(clamped); + let layer_local = unpack_layer_channels(clamped); + + // Weight index calculation + let ky_idx = u32(ky + KERNEL_RADIUS); + let kx_idx = u32(kx + KERNEL_RADIUS); + let spatial_idx = ky_idx * KERNEL_SIZE + kx_idx; + + // Accumulate: static features (8D) + for (var i: u32 = 0u; i < 8u; i++) { + let w_idx = c * IN_CHANNELS * KERNEL_SIZE * KERNEL_SIZE + + i * KERNEL_SIZE * KERNEL_SIZE + spatial_idx; + sum += weights[w_idx] * static_local[i]; + } + + // Accumulate: layer input channels (if layer_idx > 0) + let prev_channels = IN_CHANNELS - 8u; + for (var i: u32 = 0u; i < prev_channels; i++) { + let w_idx = c * IN_CHANNELS * KERNEL_SIZE * KERNEL_SIZE + + (8u + i) * KERNEL_SIZE * KERNEL_SIZE + spatial_idx; + sum += weights[w_idx] * layer_local[i]; + } + } + } + + output[c] = clamp(sum, 0.0, 1.0); // Sigmoid approximation + } + + // Pack and store + textureStore(output_tex, coord, pack_channels(output)); +} diff --git a/cnn_v2/shaders/cnn_v2_layer_template.wgsl b/cnn_v2/shaders/cnn_v2_layer_template.wgsl new file mode 100644 index 0000000..1bf6819 --- /dev/null +++ b/cnn_v2/shaders/cnn_v2_layer_template.wgsl @@ -0,0 +1,68 @@ +// CNN v2 Layer Template (placeholder for generated shaders) +// This file documents the structure - actual layers generated by export script + +// Example: Layer 0 (1×1 kernel, 8→16 channels) +// const KERNEL_SIZE: u32 = 1u; +// const IN_CHANNELS: u32 = 8u; // 7 features + bias +// const OUT_CHANNELS: u32 = 16u; +// const weights: array = array(...); + +@group(0) @binding(0) var static_features: texture_2d; +@group(0) @binding(1) var layer_input: texture_2d; // Previous layer output +@group(0) @binding(2) var output_tex: texture_storage_2d; + +fn unpack_static_features(coord: vec2) -> array { + let packed = textureLoad(static_features, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn unpack_layer_channels(coord: vec2) -> array { + let packed = textureLoad(layer_input, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn pack_channels(values: array) -> vec4 { + return vec4( + pack2x16float(vec2(values[0], values[1])), + pack2x16float(vec2(values[2], values[3])), + pack2x16float(vec2(values[4], values[5])), + pack2x16float(vec2(values[6], values[7])) + ); +} + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) { + let coord = vec2(id.xy); + let dims = textureDimensions(static_features); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) { + return; + } + + // Load static features (always available) + let static_feat = unpack_static_features(coord); + + // Convolution loop (example for generated code) + // var output: array; + // for (var c: u32 = 0u; c < OUT_CHANNELS; c++) { + // var sum: f32 = 0.0; + // for (var ky: i32 = -radius; ky <= radius; ky++) { + // for (var kx: i32 = -radius; kx <= radius; kx++) { + // let sample_coord = coord + vec2(kx, ky); + // // Load static + prev layer, multiply weights, accumulate + // } + // } + // output[c] = max(0.0, sum); // ReLU + // } + + // Placeholder output + textureStore(output_tex, coord, vec4(0u)); +} diff --git a/cnn_v2/shaders/cnn_v2_static.wgsl b/cnn_v2/shaders/cnn_v2_static.wgsl new file mode 100644 index 0000000..309e832 --- /dev/null +++ b/cnn_v2/shaders/cnn_v2_static.wgsl @@ -0,0 +1,75 @@ +// CNN v2 Static Features Compute Shader +// Generates 8D parametric features: [p0, p1, p2, p3, uv.x, uv.y, sin20_y, bias] +// p0-p3: Parametric features from specified mip level (0=mip0, 1=mip1, 2=mip2, 3=mip3) +// Note: Input image RGBD (mip0) fed separately to Layer 0 +// +// TODO: Binary format should support arbitrary layout and ordering for feature vector (7D). +// Current layout is hardcoded. Future versions should allow runtime-specified +// feature combinations (e.g., [R, G, B, dx, dy, uv_x, bias] or custom encodings). + +struct StaticFeatureParams { + mip_level: u32, + padding0: u32, + padding1: u32, + padding2: u32, +} + +@group(0) @binding(0) var input_tex: texture_2d; +@group(0) @binding(1) var input_tex_mip1: texture_2d; +@group(0) @binding(2) var input_tex_mip2: texture_2d; +@group(0) @binding(3) var depth_tex: texture_2d; +@group(0) @binding(4) var output_tex: texture_storage_2d; +@group(0) @binding(5) var params: StaticFeatureParams; +@group(0) @binding(6) var linear_sampler: sampler; + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) { + let coord = vec2(id.xy); + let dims = textureDimensions(input_tex); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) { + return; + } + + // Parametric features (p0-p3) - bilinear sample from specified mip level + // Use UV coordinates for bilinear interpolation + // Note: Use textureSampleLevel (not textureSample) in compute shaders + let uv = (vec2(coord) + 0.5) / vec2(dims); + var rgba: vec4; + if (params.mip_level == 0u) { + rgba = textureSampleLevel(input_tex, linear_sampler, uv, 0.0); + } else if (params.mip_level == 1u) { + rgba = textureSampleLevel(input_tex_mip1, linear_sampler, uv, 0.0); + } else if (params.mip_level == 2u) { + rgba = textureSampleLevel(input_tex_mip2, linear_sampler, uv, 0.0); + } else { + // Mip 3 or higher: use mip 2 as fallback + rgba = textureSampleLevel(input_tex_mip2, linear_sampler, uv, 0.0); + } + + let p0 = rgba.r; + let p1 = rgba.g; + let p2 = rgba.b; + let p3 = textureLoad(depth_tex, coord, 0).r; + + // UV coordinates (normalized [0,1], top-left origin - matches training) + let uv_x = f32(coord.x) / f32(dims.x); + let uv_y = f32(coord.y) / f32(dims.y); + + // Multi-frequency position encoding + let sin20_y = sin(20.0 * uv_y); + + // Bias dimension (always 1.0) + let bias = 1.0; + + // Pack 8×f16 into 4×u32 (rgba32uint) + // [p0, p1, p2, p3, uv_x, uv_y, sin20_y, bias] + let packed = vec4( + pack2x16float(vec2(p0, p1)), + pack2x16float(vec2(p2, p3)), + pack2x16float(vec2(uv_x, uv_y)), + pack2x16float(vec2(sin20_y, bias)) + ); + + textureStore(output_tex, coord, packed); +} diff --git a/cnn_v2/src/cnn_v2_effect.cc b/cnn_v2/src/cnn_v2_effect.cc new file mode 100644 index 0000000..60538d4 --- /dev/null +++ b/cnn_v2/src/cnn_v2_effect.cc @@ -0,0 +1,497 @@ +// CNN v2 Effect Implementation + +#include "cnn_v2_effect.h" + +#if defined(USE_TEST_ASSETS) +#include "test_assets.h" +#else +#include "generated/assets.h" +#endif + +#include "gpu/bind_group_builder.h" +#include "gpu/gpu.h" +#include "util/asset_manager.h" +#include "util/fatal_error.h" +#include + +CNNv2Effect::CNNv2Effect(const GpuContext& ctx) + : PostProcessEffect(ctx), static_pipeline_(nullptr), + static_bind_group_(nullptr), static_params_buffer_(nullptr), + static_features_tex_(nullptr), static_features_view_(nullptr), + linear_sampler_(nullptr), layer_pipeline_(nullptr), + weights_buffer_(nullptr), input_mip_tex_(nullptr), + current_input_view_(nullptr), blend_amount_(1.0f), mip_level_(0), + initialized_(false) { + std::memset(input_mip_view_, 0, sizeof(input_mip_view_)); +} + +CNNv2Effect::CNNv2Effect(const GpuContext& ctx, const CNNv2EffectParams& params) + : PostProcessEffect(ctx), static_pipeline_(nullptr), + static_bind_group_(nullptr), static_params_buffer_(nullptr), + static_features_tex_(nullptr), static_features_view_(nullptr), + linear_sampler_(nullptr), layer_pipeline_(nullptr), + weights_buffer_(nullptr), input_mip_tex_(nullptr), + current_input_view_(nullptr), blend_amount_(params.blend_amount), + mip_level_(0), initialized_(false) { + std::memset(input_mip_view_, 0, sizeof(input_mip_view_)); +} + +CNNv2Effect::~CNNv2Effect() { + cleanup(); +} + +void CNNv2Effect::init(MainSequence* demo) { + (void)demo; + if (initialized_) + return; + + load_weights(); + create_textures(); + create_pipelines(); + + initialized_ = true; +} + +void CNNv2Effect::resize(int width, int height) { + PostProcessEffect::resize(width, height); + cleanup(); + create_textures(); + create_pipelines(); +} + +void CNNv2Effect::load_weights() { + // Load binary weights asset + size_t weights_size = 0; + const uint8_t* weights_data = + (const uint8_t*)GetAsset(AssetId::ASSET_WEIGHTS_CNN_V2, &weights_size); + + if (!weights_data || weights_size < 20) { + // Weights not available - effect will skip + return; + } + + // Parse header + const uint32_t* header = (const uint32_t*)weights_data; + uint32_t magic = header[0]; + uint32_t version = header[1]; + uint32_t num_layers = header[2]; + uint32_t total_weights = header[3]; + + FATAL_CHECK(magic != 0x324e4e43, "Invalid CNN v2 weights magic\n"); // 'CNN2' + + // Support both version 1 (16-byte header) and version 2 (20-byte header with + // mip_level) + // TODO: Version 3 should include feature descriptor for arbitrary + // layout/ordering + if (version == 1) { + mip_level_ = 0; // Default for v1 + } else if (version == 2) { + mip_level_ = header[4]; + } else { + FATAL_ERROR("Unsupported CNN v2 weights version: %u\n", version); + } + + // Parse layer info (20 bytes per layer) + // Offset depends on version: v1=16 bytes (4 u32), v2=20 bytes (5 u32) + const uint32_t header_u32_count = (version == 1) ? 4 : 5; + const uint32_t* layer_data = header + header_u32_count; + for (uint32_t i = 0; i < num_layers; ++i) { + LayerInfo info; + info.kernel_size = layer_data[i * 5 + 0]; + info.in_channels = layer_data[i * 5 + 1]; + info.out_channels = layer_data[i * 5 + 2]; + info.weight_offset = layer_data[i * 5 + 3]; + info.weight_count = layer_data[i * 5 + 4]; + layer_info_.push_back(info); + } + + // Create GPU storage buffer for weights (skip header + layer info, upload + // only weights) + size_t header_size = 20; // 5 u32 + size_t layer_info_size = 20 * num_layers; // 5 u32 per layer + size_t weights_offset = header_size + layer_info_size; + size_t weights_only_size = weights_size - weights_offset; + + WGPUBufferDescriptor buffer_desc = {}; + buffer_desc.size = weights_only_size; + buffer_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst; + buffer_desc.mappedAtCreation = false; + + weights_buffer_ = wgpuDeviceCreateBuffer(ctx_.device, &buffer_desc); + + // Upload only weights (skip header + layer info) + wgpuQueueWriteBuffer(ctx_.queue, weights_buffer_, 0, + weights_data + weights_offset, weights_only_size); + + // Create uniform buffers for layer params (one per layer) + for (uint32_t i = 0; i < num_layers; ++i) { + WGPUBufferDescriptor params_desc = {}; + params_desc.size = sizeof(LayerParams); + params_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + params_desc.mappedAtCreation = false; + + WGPUBuffer buf = wgpuDeviceCreateBuffer(ctx_.device, ¶ms_desc); + layer_params_buffers_.push_back(buf); + } +} + +void CNNv2Effect::create_textures() { + // Static features texture (8×f16 packed as 4×u32) + TextureWithView static_tex = gpu_create_storage_texture_2d( + ctx_.device, width_, height_, WGPUTextureFormat_RGBA32Uint); + static_features_tex_ = static_tex.texture; + static_features_view_ = static_tex.view; + + // Input texture with mips (for multi-scale features) + TextureWithView input_mip = gpu_create_texture_2d( + ctx_.device, width_, height_, WGPUTextureFormat_RGBA8Unorm, + (WGPUTextureUsage)(WGPUTextureUsage_TextureBinding | + WGPUTextureUsage_CopyDst), + 3); + input_mip_tex_ = input_mip.texture; + + for (int i = 0; i < 3; ++i) { + input_mip_view_[i] = + gpu_create_mip_view(input_mip_tex_, WGPUTextureFormat_RGBA8Unorm, i); + } + + // Create 2 layer textures (ping-pong buffers for intermediate results) + // Each stores 8×f16 channels packed as 4×u32 + for (int i = 0; i < 2; ++i) { + TextureWithView layer = gpu_create_storage_texture_2d( + ctx_.device, width_, height_, WGPUTextureFormat_RGBA32Uint); + layer_textures_.push_back(layer.texture); + layer_views_.push_back(layer.view); + } + + // Create uniform buffer for static feature params + WGPUBufferDescriptor params_desc = {}; + params_desc.size = sizeof(StaticFeatureParams); + params_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + params_desc.mappedAtCreation = false; + static_params_buffer_ = wgpuDeviceCreateBuffer(ctx_.device, ¶ms_desc); +} + +void CNNv2Effect::create_pipelines() { + // Create linear sampler for bilinear interpolation + WGPUSamplerDescriptor sampler_desc = {}; + sampler_desc.addressModeU = WGPUAddressMode_ClampToEdge; + sampler_desc.addressModeV = WGPUAddressMode_ClampToEdge; + sampler_desc.addressModeW = WGPUAddressMode_ClampToEdge; + sampler_desc.magFilter = WGPUFilterMode_Linear; + sampler_desc.minFilter = WGPUFilterMode_Linear; + sampler_desc.mipmapFilter = WGPUMipmapFilterMode_Linear; + sampler_desc.lodMinClamp = 0.0f; + sampler_desc.lodMaxClamp = 32.0f; + sampler_desc.maxAnisotropy = 1; + + linear_sampler_ = wgpuDeviceCreateSampler(ctx_.device, &sampler_desc); + + // Static features compute pipeline + size_t shader_size = 0; + const char* static_code = + (const char*)GetAsset(AssetId::ASSET_SHADER_CNN_V2_STATIC, &shader_size); + + if (!static_code || shader_size == 0) { + // Shader not available (e.g., in test mode) - skip pipeline creation + return; + } + + WGPUShaderSourceWGSL wgsl_src = {}; + wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_src.code = str_view(static_code); + + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_src.chain; + + // Create bind group layout for static features compute + // Bindings: 0=input_tex, 1=input_mip1, 2=input_mip2, 3=depth_tex, 4=output, + // 5=params, 6=linear_sampler + WGPUBindGroupLayout static_bgl = + BindGroupLayoutBuilder() + .texture(0, WGPUShaderStage_Compute) + .texture(1, WGPUShaderStage_Compute) + .texture(2, WGPUShaderStage_Compute) + .texture(3, WGPUShaderStage_Compute) + .storage_texture(4, WGPUShaderStage_Compute, + WGPUTextureFormat_RGBA32Uint) + .uniform(5, WGPUShaderStage_Compute, sizeof(StaticFeatureParams)) + .sampler(6, WGPUShaderStage_Compute) + .build(ctx_.device); + + // Update pipeline layout + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &static_bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(ctx_.device, &pl_desc); + + // Recreate pipeline with proper layout + WGPUComputePipelineDescriptor pipeline_desc2 = {}; + pipeline_desc2.compute.module = + wgpuDeviceCreateShaderModule(ctx_.device, &shader_desc); + pipeline_desc2.compute.entryPoint = str_view("main"); + pipeline_desc2.layout = pipeline_layout; + + if (static_pipeline_) + wgpuComputePipelineRelease(static_pipeline_); + static_pipeline_ = + wgpuDeviceCreateComputePipeline(ctx_.device, &pipeline_desc2); + + wgpuShaderModuleRelease(pipeline_desc2.compute.module); + wgpuPipelineLayoutRelease(pipeline_layout); + wgpuBindGroupLayoutRelease(static_bgl); + + // CNN layer compute pipeline (storage buffer version) + if (layer_info_.empty()) + return; // No weights loaded + + size_t layer_shader_size = 0; + const char* layer_code = (const char*)GetAsset( + AssetId::ASSET_SHADER_CNN_V2_COMPUTE, &layer_shader_size); + + if (!layer_code || layer_shader_size == 0) + return; + + WGPUShaderSourceWGSL layer_wgsl = {}; + layer_wgsl.chain.sType = WGPUSType_ShaderSourceWGSL; + layer_wgsl.code = str_view(layer_code); + + WGPUShaderModuleDescriptor layer_shader_desc = {}; + layer_shader_desc.nextInChain = &layer_wgsl.chain; + + WGPUShaderModule layer_module = + wgpuDeviceCreateShaderModule(ctx_.device, &layer_shader_desc); + if (!layer_module) + return; + + // Create bind group layout for layer compute + // 0=static_features, 1=layer_input, 2=output, 3=weights, 4=params, + // 5=original_input + WGPUBindGroupLayout layer_bgl = + BindGroupLayoutBuilder() + .uint_texture(0, WGPUShaderStage_Compute) + .uint_texture(1, WGPUShaderStage_Compute) + .storage_texture(2, WGPUShaderStage_Compute, + WGPUTextureFormat_RGBA32Uint) + .storage(3, WGPUShaderStage_Compute) + .uniform(4, WGPUShaderStage_Compute, sizeof(LayerParams)) + .texture(5, WGPUShaderStage_Compute) + .build(ctx_.device); + + WGPUPipelineLayoutDescriptor layer_pl_desc = {}; + layer_pl_desc.bindGroupLayoutCount = 1; + layer_pl_desc.bindGroupLayouts = &layer_bgl; + + WGPUPipelineLayout layer_pipeline_layout = + wgpuDeviceCreatePipelineLayout(ctx_.device, &layer_pl_desc); + + WGPUComputePipelineDescriptor layer_pipeline_desc = {}; + layer_pipeline_desc.compute.module = layer_module; + layer_pipeline_desc.compute.entryPoint = str_view("main"); + layer_pipeline_desc.layout = layer_pipeline_layout; + + layer_pipeline_ = + wgpuDeviceCreateComputePipeline(ctx_.device, &layer_pipeline_desc); + + wgpuShaderModuleRelease(layer_module); + wgpuPipelineLayoutRelease(layer_pipeline_layout); + wgpuBindGroupLayoutRelease(layer_bgl); +} + +void CNNv2Effect::update_bind_group(WGPUTextureView input_view) { + if (!static_pipeline_) + return; + + // Cache input view + current_input_view_ = input_view; + + // Release old bind group + if (static_bind_group_) { + wgpuBindGroupRelease(static_bind_group_); + static_bind_group_ = nullptr; + } + + // Create bind group for static features compute (manual for storage texture + // binding) + WGPUBindGroupEntry bg_entries[7] = {}; + bg_entries[0].binding = 0; + bg_entries[0].textureView = input_view; + bg_entries[1].binding = 1; + bg_entries[1].textureView = input_mip_view_[0]; + bg_entries[2].binding = 2; + bg_entries[2].textureView = + input_mip_view_[1] ? input_mip_view_[1] : input_mip_view_[0]; + bg_entries[3].binding = 3; + bg_entries[3].textureView = input_view; + bg_entries[4].binding = 4; + bg_entries[4].textureView = static_features_view_; + bg_entries[5].binding = 5; + bg_entries[5].buffer = static_params_buffer_; + bg_entries[5].size = sizeof(StaticFeatureParams); + bg_entries[6].binding = 6; + bg_entries[6].sampler = linear_sampler_; + + WGPUBindGroupLayout layout = + wgpuComputePipelineGetBindGroupLayout(static_pipeline_, 0); + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = layout; + bg_desc.entryCount = 7; + bg_desc.entries = bg_entries; + static_bind_group_ = wgpuDeviceCreateBindGroup(ctx_.device, &bg_desc); + wgpuBindGroupLayoutRelease(layout); + + // Create layer bind groups + if (!layer_pipeline_ || layer_info_.empty()) + return; + + // Release old layer bind groups + for (auto bg : layer_bind_groups_) { + wgpuBindGroupRelease(bg); + } + layer_bind_groups_.clear(); + + // Get bind group layout from layer pipeline + WGPUBindGroupLayout layer_bgl = + wgpuComputePipelineGetBindGroupLayout(layer_pipeline_, 0); + + // Create bind group for each layer + for (size_t i = 0; i < layer_info_.size(); ++i) { + WGPUTextureView layer_input = + (i == 0) ? static_features_view_ : layer_views_[i % 2]; + + WGPUBindGroup layer_bg = + BindGroupBuilder() + .texture(0, static_features_view_) + .texture(1, layer_input) + .texture(2, layer_views_[(i + 1) % 2]) + .buffer(3, weights_buffer_, wgpuBufferGetSize(weights_buffer_)) + .buffer(4, layer_params_buffers_[i], sizeof(LayerParams)) + .texture(5, input_view) + .build(ctx_.device, layer_bgl); + + layer_bind_groups_.push_back(layer_bg); + } + + wgpuBindGroupLayoutRelease(layer_bgl); +} + +void CNNv2Effect::compute(WGPUCommandEncoder encoder, + const CommonPostProcessUniforms& uniforms) { + if (!initialized_ || !static_pipeline_ || !static_bind_group_) + return; + + float effective_blend = blend_amount_; + if (beat_modulated_) { + effective_blend = blend_amount_ * uniforms.beat_phase * beat_scale_; + } + + // Update static feature params + StaticFeatureParams static_params; + static_params.mip_level = mip_level_; + static_params.padding[0] = 0; + static_params.padding[1] = 0; + static_params.padding[2] = 0; + wgpuQueueWriteBuffer(ctx_.queue, static_params_buffer_, 0, &static_params, + sizeof(static_params)); + + // Pass 1: Compute static features + WGPUComputePassEncoder pass = + wgpuCommandEncoderBeginComputePass(encoder, nullptr); + + wgpuComputePassEncoderSetPipeline(pass, static_pipeline_); + wgpuComputePassEncoderSetBindGroup(pass, 0, static_bind_group_, 0, nullptr); + + // Dispatch workgroups (8×8 threads per group) + uint32_t workgroups_x = (width_ + 7) / 8; + uint32_t workgroups_y = (height_ + 7) / 8; + wgpuComputePassEncoderDispatchWorkgroups(pass, workgroups_x, workgroups_y, 1); + + wgpuComputePassEncoderEnd(pass); + wgpuComputePassEncoderRelease(pass); + + // Execute CNN layer passes + if (!layer_pipeline_ || layer_bind_groups_.empty()) + return; + + // Update layer params (each layer has own buffer) + for (size_t i = 0; i < layer_info_.size(); ++i) { + const LayerInfo& info = layer_info_[i]; + + LayerParams params; + params.kernel_size = info.kernel_size; + params.in_channels = info.in_channels; + params.out_channels = info.out_channels; + params.weight_offset = info.weight_offset; + params.is_output_layer = (i == layer_info_.size() - 1) ? 1 : 0; + params.blend_amount = effective_blend; + params.is_layer_0 = (i == 0) ? 1 : 0; + + wgpuQueueWriteBuffer(ctx_.queue, layer_params_buffers_[i], 0, ¶ms, + sizeof(params)); + + WGPUComputePassEncoder layer_pass = + wgpuCommandEncoderBeginComputePass(encoder, nullptr); + + wgpuComputePassEncoderSetPipeline(layer_pass, layer_pipeline_); + wgpuComputePassEncoderSetBindGroup(layer_pass, 0, layer_bind_groups_[i], 0, + nullptr); + + wgpuComputePassEncoderDispatchWorkgroups(layer_pass, workgroups_x, + workgroups_y, 1); + + wgpuComputePassEncoderEnd(layer_pass); + wgpuComputePassEncoderRelease(layer_pass); + } +} + +void CNNv2Effect::render(WGPURenderPassEncoder pass, + const CommonPostProcessUniforms& uniforms) { + (void)pass; + (void)uniforms; + // Compute-only effect, rendering is done by default composite pass +} + +void CNNv2Effect::cleanup() { + if (static_features_view_) + wgpuTextureViewRelease(static_features_view_); + if (static_features_tex_) + wgpuTextureRelease(static_features_tex_); + if (static_bind_group_) + wgpuBindGroupRelease(static_bind_group_); + if (static_params_buffer_) + wgpuBufferRelease(static_params_buffer_); + if (static_pipeline_) + wgpuComputePipelineRelease(static_pipeline_); + if (linear_sampler_) + wgpuSamplerRelease(linear_sampler_); + + if (layer_pipeline_) + wgpuComputePipelineRelease(layer_pipeline_); + if (weights_buffer_) + wgpuBufferRelease(weights_buffer_); + for (auto buf : layer_params_buffers_) + wgpuBufferRelease(buf); + layer_params_buffers_.clear(); + + for (int i = 0; i < 3; ++i) { + if (input_mip_view_[i]) + wgpuTextureViewRelease(input_mip_view_[i]); + } + if (input_mip_tex_) + wgpuTextureRelease(input_mip_tex_); + + for (auto view : layer_views_) + wgpuTextureViewRelease(view); + for (auto tex : layer_textures_) + wgpuTextureRelease(tex); + for (auto bg : layer_bind_groups_) + wgpuBindGroupRelease(bg); + + layer_views_.clear(); + layer_textures_.clear(); + layer_bind_groups_.clear(); + layer_info_.clear(); + + initialized_ = false; +} diff --git a/cnn_v2/src/cnn_v2_effect.h b/cnn_v2/src/cnn_v2_effect.h new file mode 100644 index 0000000..7960b4f --- /dev/null +++ b/cnn_v2/src/cnn_v2_effect.h @@ -0,0 +1,89 @@ +// CNN v2 Effect - Parametric Static Features +// Multi-pass post-processing with 7D feature input +// Supports per-layer kernel sizes (e.g., 1×1, 3×3, 5×5) + +#pragma once +#include "gpu/effect.h" +#include + +struct CNNv2EffectParams { + float blend_amount = 1.0f; +}; + +class CNNv2Effect : public PostProcessEffect { + public: + explicit CNNv2Effect(const GpuContext& ctx); + explicit CNNv2Effect(const GpuContext& ctx, const CNNv2EffectParams& params); + ~CNNv2Effect(); + + void init(MainSequence* demo) override; + void resize(int width, int height) override; + void compute(WGPUCommandEncoder encoder, + const CommonPostProcessUniforms& uniforms) override; + void render(WGPURenderPassEncoder pass, + const CommonPostProcessUniforms& uniforms) override; + void update_bind_group(WGPUTextureView input_view) override; + + void set_beat_modulation(bool enabled, float scale = 1.0f) { + beat_modulated_ = enabled; + beat_scale_ = scale; + } + + private: + struct LayerInfo { + uint32_t kernel_size; + uint32_t in_channels; + uint32_t out_channels; + uint32_t weight_offset; + uint32_t weight_count; + }; + + struct LayerParams { + uint32_t kernel_size; + uint32_t in_channels; + uint32_t out_channels; + uint32_t weight_offset; + uint32_t is_output_layer; + float blend_amount; + uint32_t is_layer_0; + }; + + struct StaticFeatureParams { + uint32_t mip_level; + uint32_t padding[3]; + }; + + void create_textures(); + void create_pipelines(); + void load_weights(); + void cleanup(); + + // Static features compute + WGPUComputePipeline static_pipeline_; + WGPUBindGroup static_bind_group_; + WGPUBuffer static_params_buffer_; + WGPUTexture static_features_tex_; + WGPUTextureView static_features_view_; + WGPUSampler linear_sampler_; + + // CNN layers (storage buffer architecture) + WGPUComputePipeline layer_pipeline_; // Single pipeline for all layers + WGPUBuffer weights_buffer_; // Storage buffer for weights + std::vector + layer_params_buffers_; // Uniform buffers (one per layer) + std::vector layer_info_; // Layer metadata + std::vector layer_bind_groups_; // Per-layer bind groups + std::vector layer_textures_; // Ping-pong buffers + std::vector layer_views_; + + // Input mips + WGPUTexture input_mip_tex_; + WGPUTextureView input_mip_view_[3]; + WGPUTextureView current_input_view_; + + float blend_amount_ = 1.0f; + bool beat_modulated_ = false; + float beat_scale_ = 1.0f; + uint32_t mip_level_ = 0; + bool initialized_; +}; diff --git a/cnn_v2/tools/cnn_v2_test/README.md b/cnn_v2/tools/cnn_v2_test/README.md new file mode 100644 index 0000000..d41a00f --- /dev/null +++ b/cnn_v2/tools/cnn_v2_test/README.md @@ -0,0 +1,251 @@ +# CNN v2 Testing Tool + +WebGPU-based browser tool for testing trained CNN v2 weights. + +--- + +## Features + +- Drag-drop PNG images and `.bin` weights (or click to browse) +- Real-time CNN inference with WebGPU compute shaders +- View modes: CNN output, original input, difference (×10) +- Adjustable blend amount and depth +- Data-driven pipeline (supports variable layer count) +- GPU timing display +- **Left Panel:** Weights info + kernel visualization (1px/weight, all layers) +- **Right Panel:** Layer activation viewer with 4-channel split + 4× zoom + +--- + +## Requirements + +- Browser with WebGPU support: + - Chrome/Edge 113+ (enable `chrome://flags/#enable-unsafe-webgpu` if needed) + - Safari 18+ (macOS Ventura+) +- Trained CNN v2 weights in binary format (`.bin`) +- Test images (PNG format) + +--- + +## Usage + +### 1. Open Tool + +```bash +open tools/cnn_v2_test/index.html +``` + +Or use a local server to avoid CORS: +```bash +python3 -m http.server 8000 +# Open http://localhost:8000/tools/cnn_v2_test/ +``` + +### 2. Load Data + +1. **Drop `.bin` weights** into left sidebar zone (or click to browse) +2. **Drop PNG image** anywhere in center canvas area +3. CNN runs automatically when both loaded + +### 3. Layout + +**Left Sidebar:** +- Weights drop zone (click or drag-drop `.bin` files) +- Weights info panel (layer specs, ranges, file size) +- Weights visualization (click Layer 0/1/2 buttons) + - 1 pixel per weight, all input channels horizontally + - Output channels (Out 0-3) stacked vertically + +**Center Canvas:** +- Main output view (CNN result, original, or diff) +- Keyboard: `SPACE` = original, `D` = diff (×10) + +**Right Sidebar:** +- Layer selection buttons (Static 0-3/4-7, Layer 0/1/2) +- 4 small activation views (Ch0/1/2/3) in a row +- Large zoom view below (4× magnification, follows mouse) + +**Header Controls:** +- **Blend:** Mix between original (0.0) and CNN output (1.0) +- **Depth:** Uniform depth value for all pixels (0.0–1.0) +- **View:** Current display mode + +**Footer:** +- Status: GPU timing (ms), image dimensions, view mode +- Console: Timestamped event log (file loads, errors) + +--- + +## Preparing Test Data + +### Export Weights + +```bash +# From trained checkpoint +./training/export_cnn_v2_weights.py \ + checkpoints/checkpoint_epoch_100.pth \ + --output-weights tools/cnn_v2_test/test_weights.bin +``` + +Binary format: 16-byte header + 20 bytes per layer + f16 weights (~3.2 KB for 3-layer model) + +### Test Images + +Use training images or any PNG: +```bash +# Copy test image +cp training/input/test.png tools/cnn_v2_test/ +``` + +**Note:** Grayscale images automatically converted to RGB. + +--- + +## Validation + +### Visual Comparison + +Compare browser output with C++ tool: + +```bash +# Generate C++ output +./build/cnn_test training/input/test.png /tmp/cpp_output.png + +# Load same image in browser tool +# Visually compare outputs +``` + +### GPU Timing + +Expected performance: +- 512×512: ~1-2 ms (integrated GPU) +- 1024×1024: ~3-5 ms +- 1920×1080: ~5-8 ms + +Slower than expected? Check: +- WebGPU enabled in browser +- Dedicated GPU selected (if available) +- No background tabs consuming GPU + +--- + +## Troubleshooting + +### "WebGPU not supported" + +- Update browser to latest version +- Enable WebGPU flag: `chrome://flags/#enable-unsafe-webgpu` +- Try Safari 18+ (native WebGPU on macOS) + +### "Invalid .bin file" + +- Check magic number: `hexdump -C weights.bin | head` +- Should start with: `43 4e 4e 32` ('CNN2') +- Re-export weights: `./training/export_cnn_v2_weights.py` + +### Black output / incorrect colors + +- Check blend slider (set to 1.0 for full CNN output) +- Verify training converged (loss < 0.01) +- Compare with C++ tool output + +### Shader compilation errors + +Open browser console (F12) for detailed errors. Common issues: +- Image too large (>4096×4096 not tested) +- Unsupported texture format (rare on modern GPUs) + +--- + +## Architecture + +**Pipeline:** +1. **Static Features Pass** - Generate 8D features (RGBD, UV, sin, bias) +2. **CNN Layer Passes** - Compute N layers with ping-pong textures +3. **Display Pass** - Unpack and render with view mode + +**Textures:** +- Input: RGBA8 (original image) +- Depth: R32F (uniform depth) +- Static features: RGBA32Uint (8×f16 packed) +- Layer buffers: RGBA32Uint (ping-pong) + +**Data-Driven Execution:** +- Layer count read from binary header +- Per-layer params (kernel size, channels, offsets) from binary +- Single CNN shader dispatched N times + +--- + +## Implemented Features + +**✓ Weights Metadata Panel:** +- Layer descriptions (kernel size, channels, weight count) +- Weight statistics (min/max per layer) +- File size and layer count + +**✓ Weights Visualization:** +- Per-layer kernel heatmaps (1px/weight) +- All input channels displayed horizontally +- Output channels stacked vertically +- Normalized grayscale display + +**✓ Layer Activation Viewer:** +- Static features (8D split into 0-3 and 4-7 views) +- All CNN layer outputs (Layer 0/1/2...) +- 4-channel split view (grayscale per channel) +- Mouse-driven 4× zoom view + +## TODO + +**Future Enhancements:** +- Weight distribution histograms per layer +- Activation statistics (min/max/mean overlay) +- Side-by-side diff mode (browser vs C++ output) +- Export rendered layers as PNG + +--- + +## Extensions (v2+) + +Planned enhancements: + +**Variable Feature Count:** +- Binary v2: Add `num_features` to header +- Shader: Dynamic feature array or multiple textures + +**Multi-Scale Input (Mip Levels):** +- Uncomment mip bindings in static shader +- No binary format change needed + +**8-bit Quantized Weights:** +- Binary version bump (format field already present) +- Add quantization codepath in `get_weight()` function +- 2× size reduction (~1.6 KB) + +**Pre-defined Test Images:** +- Dropdown menu with training/input/*.png +- Requires local file server + +--- + +## Size + +- HTML structure: ~2 KB +- CSS styling: ~2 KB +- JavaScript logic: ~10 KB (includes zoom + weights viz) +- Static shader: ~1 KB +- CNN shader: ~3 KB +- Display shader: ~1 KB +- Layer viz shader: ~2 KB +- Zoom shader: ~1 KB +- **Total: ~22 KB** (single file, no dependencies) + +--- + +## See Also + +- `doc/CNN_V2.md` - Architecture and design +- `doc/HOWTO.md` - Training workflows +- `training/export_cnn_v2_weights.py` - Binary format +- `src/effects/cnn_v2_effect.cc` - C++ reference implementation diff --git a/cnn_v2/tools/cnn_v2_test/index.html b/cnn_v2/tools/cnn_v2_test/index.html new file mode 100644 index 0000000..84702d5 --- /dev/null +++ b/cnn_v2/tools/cnn_v2_test/index.html @@ -0,0 +1,2014 @@ + + + + + + + CNN v2 Testing Tool + + + + +
+

CNN v2 Testing Tool

+
+ +
+
+ +
+ Drop .bin Weights or Click to Browse +
+
+
Weights Info
+
+

No weights loaded

+
+
+ +
+
+ + +
+
+
+
+
+
+ + + +
+
+ + + 1.0 +
+
+ + + 1.0 +
+ +
+ +
+
+
+
Layer Visualization
+
+

Load image + weights

+
+
+
+
+
+
+ Drop PNG/video anywhere to begin + [SPACE] Original | [D] Diff (×10) +
+
+
+ + + + diff --git a/cnn_v2/training/export_cnn_v2_shader.py b/cnn_v2/training/export_cnn_v2_shader.py new file mode 100755 index 0000000..8692a62 --- /dev/null +++ b/cnn_v2/training/export_cnn_v2_shader.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +"""CNN v2 Shader Export Script - Uniform 12D→4D Architecture + +Converts PyTorch checkpoints to WGSL compute shaders with f16 weights. +Generates one shader per layer with embedded weight arrays. + +Note: Storage buffer approach (export_cnn_v2_weights.py) is preferred for size. + This script is for debugging/testing with per-layer shaders. +""" + +import argparse +import numpy as np +import torch +from pathlib import Path + +# Path resolution for running from any directory +SCRIPT_DIR = Path(__file__).parent +PROJECT_ROOT = SCRIPT_DIR.parent.parent + + +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: + layer_idx: Layer index (0, 1, 2, ...) + 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() + weights_f16 = weights_flat.astype(np.float16) + weights_f32 = weights_f16.astype(np.float32) # WGSL stores as f32 literals + + # Format weights as WGSL array + weights_str = ",\n ".join( + ", ".join(f"{w:.6f}" for w in weights_f32[i:i+8]) + for i in range(0, len(weights_f32), 8) + ) + + radius = kernel_size // 2 + if is_output_layer: + activation = "output[c] = clamp(sum, 0.0, 1.0); // Output layer" + elif layer_idx == 0: + activation = "output[c] = clamp(sum, 0.0, 1.0); // Layer 0: clamp [0,1]" + else: + activation = "output[c] = max(0.0, sum); // Middle layers: ReLU" + + 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) +const OUT_CHANNELS: u32 = 4u; // Uniform output +const KERNEL_RADIUS: i32 = {radius}; + +// Weights quantized to float16 (stored as f32 in WGSL) +const weights: array = array( + {weights_str} +); + +@group(0) @binding(0) var static_features: texture_2d; +@group(0) @binding(1) var layer_input: texture_2d; +@group(0) @binding(2) var output_tex: texture_storage_2d; + +fn unpack_static_features(coord: vec2) -> array {{ + let packed = textureLoad(static_features, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +}} + +fn unpack_layer_channels(coord: vec2) -> vec4 {{ + let packed = textureLoad(layer_input, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + return vec4(v0.x, v0.y, v1.x, v1.y); +}} + +fn pack_channels(values: vec4) -> vec4 {{ + return vec4( + pack2x16float(vec2(values.x, values.y)), + pack2x16float(vec2(values.z, values.w)), + 0u, // Unused + 0u // Unused + ); +}} + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) {{ + let coord = vec2(id.xy); + let dims = textureDimensions(static_features); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) {{ + return; + }} + + // Load static features (always available) + let static_feat = unpack_static_features(coord); + + // Convolution: 12D input (4 prev + 8 static) → 4D output + var output: vec4 = vec4(0.0); + for (var c: u32 = 0u; c < 4u; c++) {{ + var sum: f32 = 0.0; + + for (var ky: i32 = -KERNEL_RADIUS; ky <= KERNEL_RADIUS; ky++) {{ + for (var kx: i32 = -KERNEL_RADIUS; kx <= KERNEL_RADIUS; kx++) {{ + let sample_coord = coord + vec2(kx, ky); + + // Border handling (clamp) + let clamped = vec2( + clamp(sample_coord.x, 0, i32(dims.x) - 1), + clamp(sample_coord.y, 0, i32(dims.y) - 1) + ); + + // Load features at this spatial location + let static_local = unpack_static_features(clamped); + let layer_local = unpack_layer_channels(clamped); // 4D + + // Weight index calculation + let ky_idx = u32(ky + KERNEL_RADIUS); + let kx_idx = u32(kx + KERNEL_RADIUS); + let spatial_idx = ky_idx * KERNEL_SIZE + kx_idx; + + // Accumulate: previous/input channels (4D) + for (var i: u32 = 0u; i < 4u; i++) {{ + let w_idx = c * 12u * KERNEL_SIZE * KERNEL_SIZE + + i * KERNEL_SIZE * KERNEL_SIZE + spatial_idx; + sum += weights[w_idx] * layer_local[i]; + }} + + // Accumulate: static features (8D) + for (var i: u32 = 0u; i < 8u; i++) {{ + let w_idx = c * 12u * KERNEL_SIZE * KERNEL_SIZE + + (4u + i) * KERNEL_SIZE * KERNEL_SIZE + spatial_idx; + sum += weights[w_idx] * static_local[i]; + }} + }} + }} + + {activation} + }} + + // Pack and store + textureStore(output_tex, coord, pack_channels(output)); +}} +""" + + output_path = Path(output_dir) / "cnn_v2" / f"cnn_v2_layer_{layer_idx}.wgsl" + output_path.write_text(shader_code) + print(f" → {output_path}") + + +def export_checkpoint(checkpoint_path, output_dir): + """Export PyTorch checkpoint to WGSL shaders. + + Args: + checkpoint_path: Path to .pth checkpoint + output_dir: Output directory for shaders + """ + print(f"Loading checkpoint: {checkpoint_path}") + checkpoint = torch.load(checkpoint_path, map_location='cpu') + + state_dict = checkpoint['model_state_dict'] + config = checkpoint['config'] + + 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) + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"\nExporting shaders to {output_dir}/") + + # All layers uniform: 12D→4D + for i in range(num_layers): + layer_key = f'layers.{i}.weight' + if layer_key not in state_dict: + raise ValueError(f"Missing weights for layer {i}: {layer_key}") + + layer_weights = state_dict[layer_key].detach().numpy() + is_output = (i == num_layers - 1) + + export_layer_shader( + layer_idx=i, + weights=layer_weights, + kernel_size=kernel_size, + output_dir=output_dir, + mip_level=mip_level, + is_output_layer=is_output + ) + + print(f"\nExport complete! Generated {num_layers} shader files.") + + +def main(): + parser = argparse.ArgumentParser(description='Export CNN v2 checkpoint to WGSL shaders') + parser.add_argument('checkpoint', type=str, help='Path to checkpoint .pth file') + parser.add_argument('--output-dir', type=str, default=str(PROJECT_ROOT / 'workspaces/main/shaders'), + help='Output directory for shaders') + + args = parser.parse_args() + export_checkpoint(args.checkpoint, args.output_dir) + + +if __name__ == '__main__': + main() diff --git a/cnn_v2/training/export_cnn_v2_weights.py b/cnn_v2/training/export_cnn_v2_weights.py new file mode 100755 index 0000000..d66b980 --- /dev/null +++ b/cnn_v2/training/export_cnn_v2_weights.py @@ -0,0 +1,288 @@ +#!/usr/bin/env python3 +"""CNN v2 Weight Export Script + +Converts PyTorch checkpoints to binary weight format for storage buffer. +Exports single shader template + binary weights asset. +""" + +import argparse +import numpy as np +import torch +import struct +from pathlib import Path + +# Path resolution for running from any directory +SCRIPT_DIR = Path(__file__).parent +PROJECT_ROOT = SCRIPT_DIR.parent.parent + + +def export_weights_binary(checkpoint_path, output_path, quiet=False): + """Export CNN v2 weights to binary format. + + Binary format: + Header (20 bytes): + uint32 magic ('CNN2') + uint32 version (2) + uint32 num_layers + uint32 total_weights (f16 count) + uint32 mip_level (0-3) + + LayerInfo × num_layers (20 bytes each): + uint32 kernel_size + uint32 in_channels + uint32 out_channels + uint32 weight_offset (f16 index) + uint32 weight_count + + Weights (f16 array): + float16[] all_weights + + Args: + checkpoint_path: Path to .pth checkpoint + output_path: Output .bin file path + + Returns: + config dict for shader generation + """ + if not quiet: + print(f"Loading checkpoint: {checkpoint_path}") + checkpoint = torch.load(checkpoint_path, map_location='cpu') + + state_dict = checkpoint['model_state_dict'] + config = checkpoint['config'] + + # Support both old (kernel_size) and new (kernel_sizes) format + if 'kernel_sizes' in config: + kernel_sizes = config['kernel_sizes'] + elif 'kernel_size' in config: + kernel_size = config['kernel_size'] + num_layers = config.get('num_layers', 3) + kernel_sizes = [kernel_size] * num_layers + else: + kernel_sizes = [3, 3, 3] # fallback + + num_layers = config.get('num_layers', len(kernel_sizes)) + mip_level = config.get('mip_level', 0) + + if not quiet: + 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 + layers = [] + all_weights = [] + weight_offset = 0 + + for i in range(num_layers): + layer_key = f'layers.{i}.weight' + if layer_key not in state_dict: + raise ValueError(f"Missing weights for layer {i}: {layer_key}") + + layer_weights = state_dict[layer_key].detach().numpy() + layer_flat = layer_weights.flatten() + kernel_size = kernel_sizes[i] + + layers.append({ + 'kernel_size': kernel_size, + 'in_channels': 12, # 4 (input/prev) + 8 (static) + 'out_channels': 4, # Uniform output + 'weight_offset': weight_offset, + 'weight_count': len(layer_flat) + }) + all_weights.extend(layer_flat) + weight_offset += len(layer_flat) + + if not quiet: + print(f" Layer {i}: 12D→4D, {kernel_size}×{kernel_size}, {len(layer_flat)} weights") + + # Convert to f16 + # TODO: Use 8-bit quantization for 2× size reduction + # Requires quantization-aware training (QAT) to maintain accuracy + all_weights_f16 = np.array(all_weights, dtype=np.float16) + + # Pack f16 pairs into u32 for storage buffer + # Pad to even count if needed + if len(all_weights_f16) % 2 == 1: + all_weights_f16 = np.append(all_weights_f16, np.float16(0.0)) + + # Pack pairs using numpy view + weights_u32 = all_weights_f16.view(np.uint32) + + binary_size = 20 + len(layers) * 20 + len(weights_u32) * 4 + if not quiet: + print(f"\nWeight statistics:") + print(f" Total layers: {len(layers)}") + print(f" Total weights: {len(all_weights_f16)} (f16)") + print(f" Packed: {len(weights_u32)} u32") + print(f" Binary size: {binary_size} bytes") + + # Write binary file + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + with open(output_path, 'wb') as f: + # Header (20 bytes) - version 2 with mip_level + f.write(struct.pack('<4sIIII', + b'CNN2', # magic + 2, # version (bumped to 2) + len(layers), # num_layers + len(all_weights_f16), # total_weights (f16 count) + mip_level)) # mip_level + + # Layer info (20 bytes per layer) + for layer in layers: + f.write(struct.pack('; +@group(0) @binding(1) var layer_input: texture_2d; +@group(0) @binding(2) var output_tex: texture_storage_2d; +@group(0) @binding(3) var weights: array; // Packed f16 pairs + +fn unpack_static_features(coord: vec2) -> array { + let packed = textureLoad(static_features, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + let v2 = unpack2x16float(packed.z); + let v3 = unpack2x16float(packed.w); + return array(v0.x, v0.y, v1.x, v1.y, v2.x, v2.y, v3.x, v3.y); +} + +fn unpack_layer_channels(coord: vec2) -> vec4 { + let packed = textureLoad(layer_input, coord, 0); + let v0 = unpack2x16float(packed.x); + let v1 = unpack2x16float(packed.y); + return vec4(v0.x, v0.y, v1.x, v1.y); +} + +fn pack_channels(values: vec4) -> vec4 { + return vec4( + pack2x16float(vec2(values.x, values.y)), + pack2x16float(vec2(values.z, values.w)), + 0u, // Unused + 0u // Unused + ); +} + +fn get_weight(idx: u32) -> f32 { + let pair_idx = idx / 2u; + let packed = weights[8u + pair_idx]; // Skip header (32 bytes = 8 u32) + let unpacked = unpack2x16float(packed); + return select(unpacked.y, unpacked.x, (idx & 1u) == 0u); +} + +@compute @workgroup_size(8, 8) +fn main(@builtin(global_invocation_id) id: vec3) { + let coord = vec2(id.xy); + let dims = textureDimensions(static_features); + + if (coord.x >= i32(dims.x) || coord.y >= i32(dims.y)) { + return; + } + + // Read header + let header_packed = weights[0]; // magic + version + let counts_packed = weights[1]; // num_layers + total_weights + let num_layers = counts_packed & 0xFFFFu; + + // Load static features + let static_feat = unpack_static_features(coord); + + // Process each layer (hardcoded for 3 layers for now) + // TODO: Dynamic layer loop when needed + + // Example for layer 0 - expand to full multi-layer when tested + let layer_info_offset = 2u; // After header + let layer0_info_base = layer_info_offset; + + // Read layer 0 info (5 u32 values = 20 bytes) + let kernel_size = weights[layer0_info_base]; + let in_channels = weights[layer0_info_base + 1u]; + let out_channels = weights[layer0_info_base + 2u]; + let weight_offset = weights[layer0_info_base + 3u]; + + // Convolution: 12D input (4 prev + 8 static) → 4D output + var output: vec4 = vec4(0.0); + for (var c: u32 = 0u; c < 4u; c++) { + output[c] = 0.0; // TODO: Actual convolution + } + + textureStore(output_tex, coord, pack_channels(output)); +} +""" + + output_path = Path(output_dir) / "cnn_v2" / "cnn_v2_compute.wgsl" + output_path.write_text(shader_code) + print(f" → {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description='Export CNN v2 weights to binary format') + parser.add_argument('checkpoint', type=str, help='Path to checkpoint .pth file') + parser.add_argument('--output-weights', type=str, default=str(PROJECT_ROOT / 'workspaces/main/weights/cnn_v2_weights.bin'), + help='Output binary weights file') + parser.add_argument('--output-shader', type=str, default=str(PROJECT_ROOT / 'workspaces/main/shaders'), + help='Output directory for shader template') + parser.add_argument('--quiet', action='store_true', + help='Suppress detailed output') + + args = parser.parse_args() + + if not args.quiet: + print("=== CNN v2 Weight Export ===\n") + config = export_weights_binary(args.checkpoint, args.output_weights, quiet=args.quiet) + if not args.quiet: + print() + # Shader is manually maintained in cnn_v2_compute.wgsl + # export_shader_template(config, args.output_shader) + print("\nExport complete!") + + +if __name__ == '__main__': + main() diff --git a/cnn_v2/training/gen_identity_weights.py b/cnn_v2/training/gen_identity_weights.py new file mode 100755 index 0000000..08eecc6 --- /dev/null +++ b/cnn_v2/training/gen_identity_weights.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +"""Generate Identity CNN v2 Weights + +Creates trivial .bin with 1 layer, 1×1 kernel, identity passthrough. +Output Ch{0,1,2,3} = Input Ch{0,1,2,3} (ignores static features). + +With --mix: Output Ch{i} = 0.5*prev[i] + 0.5*static_p{4+i} + (50-50 blend of prev layer with uv_x, uv_y, sin20_y, bias) + +With --p47: Output Ch{i} = static p{4+i} (uv_x, uv_y, sin20_y, bias) + (p4/uv_x→ch0, p5/uv_y→ch1, p6/sin20_y→ch2, p7/bias→ch3) + +Usage: + ./training/gen_identity_weights.py [output.bin] + ./training/gen_identity_weights.py --mix [output.bin] + ./training/gen_identity_weights.py --p47 [output.bin] +""" + +import argparse +import numpy as np +import struct +from pathlib import Path + +# Path resolution for running from any directory +SCRIPT_DIR = Path(__file__).parent +PROJECT_ROOT = SCRIPT_DIR.parent.parent + + +def generate_identity_weights(output_path, kernel_size=1, mip_level=0, mix=False, p47=False): + """Generate identity weights: output = input (ignores static features). + + If mix=True, 50-50 blend: 0.5*p0+0.5*p4, 0.5*p1+0.5*p5, etc (avoids overflow). + If p47=True, transfers static p4-p7 (uv_x, uv_y, sin20_y, bias) to output channels. + + Input channel layout: [0-3: prev layer, 4-11: static (p0-p7)] + Static features: p0-p3 (RGB+D), p4 (uv_x), p5 (uv_y), p6 (sin20_y), p7 (bias) + + Binary format: + Header (20 bytes): + uint32 magic ('CNN2') + uint32 version (2) + uint32 num_layers (1) + uint32 total_weights (f16 count) + uint32 mip_level + + LayerInfo (20 bytes): + uint32 kernel_size + uint32 in_channels (12) + uint32 out_channels (4) + uint32 weight_offset (0) + uint32 weight_count + + Weights (u32 packed f16): + Identity matrix for first 4 input channels + Zeros for static features (channels 4-11) OR + Mix matrix (p0+p4, p1+p5, p2+p6, p3+p7) if mix=True + """ + # Identity: 4 output channels, 12 input channels + # Weight shape: [out_ch, in_ch, kernel_h, kernel_w] + in_channels = 12 # 4 input + 8 static + out_channels = 4 + + # Identity matrix: diagonal 1.0 for first 4 channels, 0.0 for rest + weights = np.zeros((out_channels, in_channels, kernel_size, kernel_size), dtype=np.float32) + + # Center position for kernel + center = kernel_size // 2 + + if p47: + # p47 mode: p4→ch0, p5→ch1, p6→ch2, p7→ch3 (static features only) + # Input channels: [0-3: prev layer, 4-11: static features (p0-p7)] + # p4-p7 are at input channels 8-11 + for i in range(out_channels): + weights[i, i + 8, center, center] = 1.0 + elif mix: + # Mix mode: 50-50 blend (p0+p4, p1+p5, p2+p6, p3+p7) + # p0-p3 are at channels 0-3 (prev layer), p4-p7 at channels 8-11 (static) + for i in range(out_channels): + weights[i, i, center, center] = 0.5 # 0.5*p{i} (prev layer) + weights[i, i + 8, center, center] = 0.5 # 0.5*p{i+4} (static) + else: + # Identity: output ch i = input ch i + for i in range(out_channels): + weights[i, i, center, center] = 1.0 + + # Flatten + weights_flat = weights.flatten() + weight_count = len(weights_flat) + + mode_name = 'p47' if p47 else ('mix' if mix else 'identity') + print(f"Generating {mode_name} weights:") + print(f" Kernel size: {kernel_size}×{kernel_size}") + print(f" Channels: 12D→4D") + print(f" Weights: {weight_count}") + print(f" Mip level: {mip_level}") + if mix: + print(f" Mode: 0.5*prev[i] + 0.5*static_p{{4+i}} (blend with uv/sin/bias)") + elif p47: + print(f" Mode: p4→ch0, p5→ch1, p6→ch2, p7→ch3") + + # Convert to f16 + weights_f16 = np.array(weights_flat, dtype=np.float16) + + # Pad to even count + if len(weights_f16) % 2 == 1: + weights_f16 = np.append(weights_f16, np.float16(0.0)) + + # Pack f16 pairs into u32 + weights_u32 = weights_f16.view(np.uint32) + + print(f" Packed: {len(weights_u32)} u32") + print(f" Binary size: {20 + 20 + len(weights_u32) * 4} bytes") + + # Write binary + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + with open(output_path, 'wb') as f: + # Header (20 bytes) + f.write(struct.pack('<4sIIII', + b'CNN2', # magic + 2, # version + 1, # num_layers + len(weights_f16), # total_weights + mip_level)) # mip_level + + # Layer info (20 bytes) + f.write(struct.pack(' 0: + # Downsample to mip level + mip_rgb = rgb.copy() + for _ in range(mip_level): + mip_rgb = cv2.pyrDown(mip_rgb) + # Upsample back to original size + for _ in range(mip_level): + mip_rgb = cv2.pyrUp(mip_rgb) + # Crop/pad to exact original size if needed + if mip_rgb.shape[:2] != (h, w): + mip_rgb = cv2.resize(mip_rgb, (w, h), interpolation=cv2.INTER_LINEAR) + else: + mip_rgb = rgb + + # Parametric features (p0-p3) from mip level + p0 = mip_rgb[:, :, 0].astype(np.float32) + p1 = mip_rgb[:, :, 1].astype(np.float32) + p2 = mip_rgb[:, :, 2].astype(np.float32) + p3 = depth.astype(np.float32) if depth is not None else np.ones((h, w), dtype=np.float32) # Default 1.0 = far plane + + # UV coordinates (normalized [0, 1]) + uv_x = np.linspace(0, 1, w)[None, :].repeat(h, axis=0).astype(np.float32) + uv_y = np.linspace(0, 1, h)[:, None].repeat(w, axis=1).astype(np.float32) + + # Multi-frequency position encoding + sin20_y = np.sin(20.0 * uv_y).astype(np.float32) + + # Bias dimension (always 1.0) - replaces Conv2d bias parameter + bias = np.ones((h, w), dtype=np.float32) + + # Stack: [p0, p1, p2, p3, uv.x, uv.y, sin20_y, bias] + features = np.stack([p0, p1, p2, p3, uv_x, uv_y, sin20_y, bias], axis=-1) + return features + + +class CNNv2(nn.Module): + """CNN v2 - Uniform 12D→4D Architecture + + All layers: input RGBD (4D) + static (8D) = 12D → 4 channels + Per-layer kernel sizes supported (e.g., [1, 3, 5]) + Uses bias=False (bias integrated in static features as 1.0) + + TODO: Add quantization-aware training (QAT) for 8-bit weights + - Use torch.quantization.QuantStub/DeQuantStub + - Train with fake quantization to adapt to 8-bit precision + - Target: ~1.3 KB weights (vs 2.6 KB with f16) + """ + + def __init__(self, kernel_sizes, num_layers=3): + super().__init__() + if isinstance(kernel_sizes, int): + kernel_sizes = [kernel_sizes] * num_layers + assert len(kernel_sizes) == num_layers, "kernel_sizes must match num_layers" + + self.kernel_sizes = kernel_sizes + self.num_layers = num_layers + self.layers = nn.ModuleList() + + # All layers: 12D input (4 RGBD + 8 static) → 4D output + for kernel_size in kernel_sizes: + self.layers.append( + nn.Conv2d(12, 4, kernel_size=kernel_size, + padding=kernel_size//2, bias=False) + ) + + def forward(self, input_rgbd, static_features): + """Forward pass with uniform 12D→4D layers. + + Args: + input_rgbd: (B, 4, H, W) input image RGBD (mip 0) + static_features: (B, 8, H, W) static features + + Returns: + (B, 4, H, W) RGBA output [0, 1] + """ + # Layer 0: input RGBD (4D) + static (8D) = 12D + x = torch.cat([input_rgbd, static_features], dim=1) + x = self.layers[0](x) + x = torch.sigmoid(x) # Soft [0,1] for layer 0 + + # Layer 1+: previous (4D) + static (8D) = 12D + for i in range(1, self.num_layers): + x_input = torch.cat([x, static_features], dim=1) + x = self.layers[i](x_input) + if i < self.num_layers - 1: + x = F.relu(x) + else: + x = torch.sigmoid(x) # Soft [0,1] for final layer + + return x + + +class PatchDataset(Dataset): + """Patch-based dataset extracting salient regions from images.""" + + def __init__(self, input_dir, target_dir, patch_size=32, patches_per_image=64, + detector='harris', mip_level=0): + self.input_paths = sorted(Path(input_dir).glob("*.png")) + self.target_paths = sorted(Path(target_dir).glob("*.png")) + self.patch_size = patch_size + self.patches_per_image = patches_per_image + self.detector = detector + self.mip_level = mip_level + + assert len(self.input_paths) == len(self.target_paths), \ + f"Mismatch: {len(self.input_paths)} inputs vs {len(self.target_paths)} targets" + + print(f"Found {len(self.input_paths)} image pairs") + print(f"Extracting {patches_per_image} patches per image using {detector} detector") + print(f"Total patches: {len(self.input_paths) * patches_per_image}") + + def __len__(self): + return len(self.input_paths) * self.patches_per_image + + def _detect_salient_points(self, img_array): + """Detect salient points on original image. + + TODO: Add random sampling to training vectors + - In addition to salient points, incorporate randomly-located samples + - Default: 10% random samples, 90% salient points + - Prevents overfitting to only high-gradient regions + - Improves generalization across entire image + - Configurable via --random-sample-percent parameter + """ + gray = cv2.cvtColor((img_array * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY) + h, w = gray.shape + half_patch = self.patch_size // 2 + + corners = None + if self.detector == 'harris': + corners = cv2.goodFeaturesToTrack(gray, self.patches_per_image * 2, + qualityLevel=0.01, minDistance=half_patch) + elif self.detector == 'fast': + fast = cv2.FastFeatureDetector_create(threshold=20) + keypoints = fast.detect(gray, None) + corners = np.array([[kp.pt[0], kp.pt[1]] for kp in keypoints[:self.patches_per_image * 2]]) + corners = corners.reshape(-1, 1, 2) if len(corners) > 0 else None + elif self.detector == 'shi-tomasi': + corners = cv2.goodFeaturesToTrack(gray, self.patches_per_image * 2, + qualityLevel=0.01, minDistance=half_patch, + useHarrisDetector=False) + elif self.detector == 'gradient': + grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3) + grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3) + gradient_mag = np.sqrt(grad_x**2 + grad_y**2) + threshold = np.percentile(gradient_mag, 95) + y_coords, x_coords = np.where(gradient_mag > threshold) + + if len(x_coords) > self.patches_per_image * 2: + indices = np.random.choice(len(x_coords), self.patches_per_image * 2, replace=False) + x_coords = x_coords[indices] + y_coords = y_coords[indices] + + corners = np.array([[x, y] for x, y in zip(x_coords, y_coords)]) + corners = corners.reshape(-1, 1, 2) if len(corners) > 0 else None + + # Fallback to random if no corners found + if corners is None or len(corners) == 0: + x_coords = np.random.randint(half_patch, w - half_patch, self.patches_per_image) + y_coords = np.random.randint(half_patch, h - half_patch, self.patches_per_image) + corners = np.array([[x, y] for x, y in zip(x_coords, y_coords)]) + corners = corners.reshape(-1, 1, 2) + + # Filter valid corners + valid_corners = [] + for corner in corners: + x, y = int(corner[0][0]), int(corner[0][1]) + if half_patch <= x < w - half_patch and half_patch <= y < h - half_patch: + valid_corners.append((x, y)) + if len(valid_corners) >= self.patches_per_image: + break + + # Fill with random if not enough + while len(valid_corners) < self.patches_per_image: + x = np.random.randint(half_patch, w - half_patch) + y = np.random.randint(half_patch, h - half_patch) + valid_corners.append((x, y)) + + return valid_corners + + def __getitem__(self, idx): + img_idx = idx // self.patches_per_image + patch_idx = idx % self.patches_per_image + + # Load original images (no resize) + input_img = np.array(Image.open(self.input_paths[img_idx]).convert('RGB')) / 255.0 + target_pil = Image.open(self.target_paths[img_idx]) + target_img = np.array(target_pil.convert('RGBA')) / 255.0 # Preserve alpha + + # Detect salient points on original image (use RGB only) + salient_points = self._detect_salient_points(input_img) + cx, cy = salient_points[patch_idx] + + # Extract patch + half_patch = self.patch_size // 2 + y1, y2 = cy - half_patch, cy + half_patch + x1, x2 = cx - half_patch, cx + half_patch + + input_patch = input_img[y1:y2, x1:x2] + target_patch = target_img[y1:y2, x1:x2] # RGBA + + # Extract depth from target alpha channel (or default to 1.0) + depth = target_patch[:, :, 3] if target_patch.shape[2] == 4 else None + + # Compute static features for patch + static_feat = compute_static_features(input_patch.astype(np.float32), depth=depth, mip_level=self.mip_level) + + # Input RGBD (mip 0) - add depth channel + input_rgbd = np.concatenate([input_patch, np.zeros((self.patch_size, self.patch_size, 1))], axis=-1) + + # Convert to tensors (C, H, W) + input_rgbd = torch.from_numpy(input_rgbd.astype(np.float32)).permute(2, 0, 1) + static_feat = torch.from_numpy(static_feat).permute(2, 0, 1) + target = torch.from_numpy(target_patch.astype(np.float32)).permute(2, 0, 1) # RGBA from image + + return input_rgbd, static_feat, target + + +class ImagePairDataset(Dataset): + """Dataset of input/target image pairs (full-image mode).""" + + def __init__(self, input_dir, target_dir, target_size=(256, 256), mip_level=0): + self.input_paths = sorted(Path(input_dir).glob("*.png")) + self.target_paths = sorted(Path(target_dir).glob("*.png")) + self.target_size = target_size + self.mip_level = mip_level + assert len(self.input_paths) == len(self.target_paths), \ + f"Mismatch: {len(self.input_paths)} inputs vs {len(self.target_paths)} targets" + + def __len__(self): + return len(self.input_paths) + + def __getitem__(self, idx): + # Load and resize images to fixed size + input_pil = Image.open(self.input_paths[idx]).convert('RGB') + target_pil = Image.open(self.target_paths[idx]) + + # Resize to target size + input_pil = input_pil.resize(self.target_size, Image.LANCZOS) + target_pil = target_pil.resize(self.target_size, Image.LANCZOS) + + input_img = np.array(input_pil) / 255.0 + target_img = np.array(target_pil.convert('RGBA')) / 255.0 # Preserve alpha + + # Extract depth from target alpha channel (or default to 1.0) + depth = target_img[:, :, 3] if target_img.shape[2] == 4 else None + + # Compute static features + static_feat = compute_static_features(input_img.astype(np.float32), depth=depth, mip_level=self.mip_level) + + # Input RGBD (mip 0) - add depth channel + h, w = input_img.shape[:2] + input_rgbd = np.concatenate([input_img, np.zeros((h, w, 1))], axis=-1) + + # Convert to tensors (C, H, W) + input_rgbd = torch.from_numpy(input_rgbd.astype(np.float32)).permute(2, 0, 1) + static_feat = torch.from_numpy(static_feat).permute(2, 0, 1) + target = torch.from_numpy(target_img.astype(np.float32)).permute(2, 0, 1) # RGBA from image + + return input_rgbd, static_feat, target + + +def train(args): + """Train CNN v2 model.""" + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + print(f"Training on {device}") + + # Create dataset (patch-based or full-image) + if args.full_image: + print(f"Mode: Full-image (resized to {args.image_size}x{args.image_size})") + target_size = (args.image_size, args.image_size) + dataset = ImagePairDataset(args.input, args.target, target_size=target_size, mip_level=args.mip_level) + dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True) + else: + print(f"Mode: Patch-based ({args.patch_size}x{args.patch_size} patches)") + dataset = PatchDataset(args.input, args.target, + patch_size=args.patch_size, + patches_per_image=args.patches_per_image, + detector=args.detector, + mip_level=args.mip_level) + dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True) + + # Parse kernel sizes + kernel_sizes = [int(k) for k in args.kernel_sizes.split(',')] + if len(kernel_sizes) == 1: + kernel_sizes = kernel_sizes * args.num_layers + else: + # When multiple kernel sizes provided, derive num_layers from list length + args.num_layers = len(kernel_sizes) + + # Create model + model = CNNv2(kernel_sizes=kernel_sizes, num_layers=args.num_layers).to(device) + total_params = sum(p.numel() for p in model.parameters()) + kernel_desc = ','.join(map(str, kernel_sizes)) + print(f"Model: {args.num_layers} layers, kernel sizes [{kernel_desc}], {total_params} weights") + print(f"Using mip level {args.mip_level} for p0-p3 features") + + # Optimizer and loss + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) + criterion = nn.MSELoss() + + # Training loop + print(f"\nTraining for {args.epochs} epochs...") + start_time = time.time() + + for epoch in range(1, args.epochs + 1): + model.train() + epoch_loss = 0.0 + + for input_rgbd, static_feat, target in dataloader: + input_rgbd = input_rgbd.to(device) + static_feat = static_feat.to(device) + target = target.to(device) + + optimizer.zero_grad() + output = model(input_rgbd, static_feat) + + # Compute loss (grayscale or RGBA) + if args.grayscale_loss: + # Convert RGBA to grayscale: Y = 0.299*R + 0.587*G + 0.114*B + output_gray = 0.299 * output[:, 0:1] + 0.587 * output[:, 1:2] + 0.114 * output[:, 2:3] + target_gray = 0.299 * target[:, 0:1] + 0.587 * target[:, 1:2] + 0.114 * target[:, 2:3] + loss = criterion(output_gray, target_gray) + else: + loss = criterion(output, target) + + loss.backward() + optimizer.step() + + epoch_loss += loss.item() + + avg_loss = epoch_loss / len(dataloader) + + # Print loss at every epoch (overwrite line with \r) + elapsed = time.time() - start_time + print(f"\rEpoch {epoch:4d}/{args.epochs} | Loss: {avg_loss:.6f} | Time: {elapsed:.1f}s", end='', flush=True) + + # Save checkpoint + if args.checkpoint_every > 0 and epoch % args.checkpoint_every == 0: + print() # Newline before checkpoint message + checkpoint_path = Path(args.checkpoint_dir) / f"checkpoint_epoch_{epoch}.pth" + checkpoint_path.parent.mkdir(parents=True, exist_ok=True) + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'loss': avg_loss, + 'config': { + 'kernel_sizes': kernel_sizes, + 'num_layers': args.num_layers, + 'mip_level': args.mip_level, + 'grayscale_loss': args.grayscale_loss, + 'features': ['p0', 'p1', 'p2', 'p3', 'uv.x', 'uv.y', 'sin20_y', 'bias'] + } + }, checkpoint_path) + print(f" → Saved checkpoint: {checkpoint_path}") + + # Always save final checkpoint + print() # Newline after training + final_checkpoint = Path(args.checkpoint_dir) / f"checkpoint_epoch_{args.epochs}.pth" + final_checkpoint.parent.mkdir(parents=True, exist_ok=True) + torch.save({ + 'epoch': args.epochs, + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'loss': avg_loss, + 'config': { + 'kernel_sizes': kernel_sizes, + 'num_layers': args.num_layers, + 'mip_level': args.mip_level, + 'grayscale_loss': args.grayscale_loss, + 'features': ['p0', 'p1', 'p2', 'p3', 'uv.x', 'uv.y', 'sin20_y', 'bias'] + } + }, final_checkpoint) + print(f" → Saved final checkpoint: {final_checkpoint}") + + print(f"\nTraining complete! Total time: {time.time() - start_time:.1f}s") + return model + + +def main(): + parser = argparse.ArgumentParser(description='Train CNN v2 with parametric static features') + parser.add_argument('--input', type=str, required=True, help='Input images directory') + parser.add_argument('--target', type=str, required=True, help='Target images directory') + + # Training mode + parser.add_argument('--full-image', action='store_true', + help='Use full-image mode (resize all images)') + parser.add_argument('--image-size', type=int, default=256, + help='Full-image mode: resize to this size (default: 256)') + + # Patch-based mode (default) + parser.add_argument('--patch-size', type=int, default=32, + help='Patch mode: patch size (default: 32)') + parser.add_argument('--patches-per-image', type=int, default=64, + help='Patch mode: patches per image (default: 64)') + parser.add_argument('--detector', type=str, default='harris', + choices=['harris', 'fast', 'shi-tomasi', 'gradient'], + help='Patch mode: salient point detector (default: harris)') + # TODO: Add --random-sample-percent parameter (default: 10) + # Mix salient points with random samples for better generalization + + # Model architecture + parser.add_argument('--kernel-sizes', type=str, default='3', + help='Comma-separated kernel sizes per layer (e.g., "3,5,3"), single value replicates (default: 3)') + parser.add_argument('--num-layers', type=int, default=3, + help='Number of CNN layers (default: 3)') + parser.add_argument('--mip-level', type=int, default=0, choices=[0, 1, 2, 3], + help='Mip level for p0-p3 features: 0=original, 1=half, 2=quarter, 3=eighth (default: 0)') + + # Training parameters + parser.add_argument('--epochs', type=int, default=5000, help='Training epochs') + parser.add_argument('--batch-size', type=int, default=16, help='Batch size') + parser.add_argument('--lr', type=float, default=1e-3, help='Learning rate') + parser.add_argument('--grayscale-loss', action='store_true', + help='Compute loss on grayscale (Y = 0.299*R + 0.587*G + 0.114*B) instead of RGBA') + parser.add_argument('--checkpoint-dir', type=str, default='checkpoints', + help='Checkpoint directory') + parser.add_argument('--checkpoint-every', type=int, default=1000, + help='Save checkpoint every N epochs (0 = disable)') + + args = parser.parse_args() + train(args) + + +if __name__ == '__main__': + main() -- cgit v1.2.3