summaryrefslogtreecommitdiff
path: root/training/gen_identity_weights.py
diff options
context:
space:
mode:
Diffstat (limited to 'training/gen_identity_weights.py')
-rwxr-xr-xtraining/gen_identity_weights.py134
1 files changed, 134 insertions, 0 deletions
diff --git a/training/gen_identity_weights.py b/training/gen_identity_weights.py
new file mode 100755
index 0000000..a84ea87
--- /dev/null
+++ b/training/gen_identity_weights.py
@@ -0,0 +1,134 @@
+#!/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).
+
+Usage:
+ ./training/gen_identity_weights.py [output.bin]
+"""
+
+import argparse
+import numpy as np
+import struct
+from pathlib import Path
+
+
+def generate_identity_weights(output_path, kernel_size=1, mip_level=0):
+ """Generate identity weights: output = input (ignores static features).
+
+ 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)
+ """
+ # 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
+
+ # Set diagonal to 1.0 (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)
+
+ print(f"Generating identity 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}")
+
+ # 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('<IIIII',
+ kernel_size, # kernel_size
+ in_channels, # in_channels
+ out_channels, # out_channels
+ 0, # weight_offset
+ weight_count)) # weight_count
+
+ # Weights (u32 packed f16)
+ f.write(weights_u32.tobytes())
+
+ print(f" → {output_path}")
+
+ # Verify
+ print("\nVerification:")
+ with open(output_path, 'rb') as f:
+ data = f.read()
+ magic, version, num_layers, total_weights, mip = struct.unpack('<4sIIII', data[:20])
+ print(f" Magic: {magic}")
+ print(f" Version: {version}")
+ print(f" Layers: {num_layers}")
+ print(f" Total weights: {total_weights}")
+ print(f" Mip level: {mip}")
+ print(f" File size: {len(data)} bytes")
+
+
+def main():
+ parser = argparse.ArgumentParser(description='Generate identity CNN v2 weights')
+ parser.add_argument('output', type=str, nargs='?',
+ default='workspaces/main/weights/cnn_v2_identity.bin',
+ help='Output .bin file path')
+ parser.add_argument('--kernel-size', type=int, default=1,
+ help='Kernel size (default: 1×1)')
+ parser.add_argument('--mip-level', type=int, default=0,
+ help='Mip level for p0-p3 features (default: 0)')
+
+ args = parser.parse_args()
+
+ print("=== Identity Weight Generator ===\n")
+ generate_identity_weights(args.output, args.kernel_size, args.mip_level)
+ print("\nDone!")
+
+
+if __name__ == '__main__':
+ main()