diff options
Diffstat (limited to 'training/export_cnn_v2_weights.py')
| -rwxr-xr-x | training/export_cnn_v2_weights.py | 16 |
1 files changed, 9 insertions, 7 deletions
diff --git a/training/export_cnn_v2_weights.py b/training/export_cnn_v2_weights.py index 9e9e352..1086516 100755 --- a/training/export_cnn_v2_weights.py +++ b/training/export_cnn_v2_weights.py @@ -16,11 +16,12 @@ def export_weights_binary(checkpoint_path, output_path): """Export CNN v2 weights to binary format. Binary format: - Header (16 bytes): + Header (20 bytes): uint32 magic ('CNN2') - uint32 version (1) + 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 @@ -107,19 +108,20 @@ def export_weights_binary(checkpoint_path, output_path): 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: {16 + len(layers) * 20 + len(weights_u32) * 4} bytes") + print(f" Binary size: {20 + len(layers) * 20 + len(weights_u32) * 4} 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 (16 bytes) - f.write(struct.pack('<4sIII', + # Header (20 bytes) - version 2 with mip_level + f.write(struct.pack('<4sIIII', b'CNN2', # magic - 1, # version + 2, # version (bumped to 2) len(layers), # num_layers - len(all_weights_f16))) # total_weights (f16 count) + len(all_weights_f16), # total_weights (f16 count) + mip_level)) # mip_level # Layer info (20 bytes per layer) for layer in layers: |