{ "model": "setup_53_2d", "task": "mito_affinities", "architecture": "Unet3dIso (2D-collapsed)", "checkpoint": "student_60000", "params": 1927907, "channels": [ "mito_aff_z", "mito_aff_y", "mito_aff_x" ], "input_voxel_size_nm": [ 16, 16, 16 ], "output_voxel_size_nm": [ 16, 16, 16 ], "output": "raw logits (apply sigmoid / threshold at 0)", "input_norm": [ { "name": "MinMaxNormalizer", "min_value": 0, "max_value": 255 }, { "name": "LambdaNormalizer", "expression": "2*x-1" } ], "io": "input 'input' float32 [1,1,1,y,x] (one z-slice); output 'output' float32 [1,C,1,y',x']", "kind": "2D slice student distilled from the 3D model (train_2d_distill.py)", "step": 60000, "distilled_from": "setup_53", "distilled_from_repo": "cellmap/mito-aff-unetiso-setup-53", "variants": [ { "model": "setup_53_2d", "variant": "fp32", "input_size": [ 1, 300, 300 ], "output_size": [ 1, 212, 212 ], "file": "setup_53_mito_affinities_16nm_2d_in300_fp32.onnx", "file_mb": 7.75, "max_abs": 2.288818359375e-05, "flip_frac": 0.0 }, { "model": "setup_53_2d", "variant": "fp16", "input_size": [ 1, 300, 300 ], "output_size": [ 1, 212, 212 ], "file": "setup_53_mito_affinities_16nm_2d_in300_fp16.onnx", "file_mb": 3.9, "max_abs": 0.10926437377929688, "flip_frac": 4.4499822000712e-05 }, { "model": "setup_53_2d", "variant": "fp32", "input_size": [ 1, 556, 556 ], "output_size": [ 1, 468, 468 ], "file": "setup_53_mito_affinities_16nm_2d_in556_fp32.onnx", "file_mb": 7.75, "max_abs": 2.09808349609375e-05, "flip_frac": 0.0 }, { "model": "setup_53_2d", "variant": "fp16", "input_size": [ 1, 556, 556 ], "output_size": [ 1, 468, 468 ], "file": "setup_53_mito_affinities_16nm_2d_in556_fp16.onnx", "file_mb": 3.9, "max_abs": 0.10926437377929688, "flip_frac": 3.652567755131858e-05 }, { "model": "setup_53_2d", "variant": "fp32", "input_size": [ 1, 1068, 1068 ], "output_size": [ 1, 980, 980 ], "file": "setup_53_mito_affinities_16nm_2d_in1068_fp32.onnx", "file_mb": 7.75, "max_abs": 2.574920654296875e-05, "flip_frac": 0.0 }, { "model": "setup_53_2d", "variant": "fp16", "input_size": [ 1, 1068, 1068 ], "output_size": [ 1, 980, 980 ], "file": "setup_53_mito_affinities_16nm_2d_in1068_fp16.onnx", "file_mb": 3.9, "max_abs": 0.10926437377929688, "flip_frac": 4.859086491739553e-05 } ], "val_f1": { "metric": "affinity F1 (mean of 3 channels), 6 GT val crops, centered 200^3, threshold logit>0", "student": { "crop495": 0.8164, "crop496": 0.796, "crop1017": 0.9015, "crop528": 0.9217, "crop1091": 0.9166, "crop1081": 0.9481 }, "student_mean": 0.8834, "teacher_3d": { "crop495": 0.955, "crop496": 0.928, "crop1017": 0.966, "crop528": 0.9746, "crop1091": 0.9517, "crop1081": 0.9456 }, "teacher_3d_mean": 0.9535 }, "torch_version": "2.9.1", "exporter": "torch.onnx.export (TorchScript, dynamo=False)", "opset": 18, "finetune": { "head": { "final_layer": "final_conv.c", "onnx_node": "/final_conv/c/Conv", "features": "/final_conv/Gather_output_0", "weight": "final_conv.c.weight", "bias": "final_conv.c.bias", "in_channels": 32, "out_channels": 3, "activation": "none (logits)", "files": [ "setup_53_mito_affinities_16nm_2d_in1068_fp16.onnx", "setup_53_mito_affinities_16nm_2d_in1068_fp32.onnx", "setup_53_mito_affinities_16nm_2d_in300_fp16.onnx", "setup_53_mito_affinities_16nm_2d_in300_fp32.onnx", "setup_53_mito_affinities_16nm_2d_in556_fp16.onnx", "setup_53_mito_affinities_16nm_2d_in556_fp32.onnx" ], "initializers": "finetune/initializers.json", "loss": "masked BCE on logits", "note": "train the final 1x1 conv on frozen features" }, "full": { "library": "tfjs", "files": [ "finetune/train_graph.json", "finetune/cellmap_tfjs_train.js" ], "weight_map": "identity", "initializers": "finetune/initializers.json", "write_back": [ "setup_53_mito_affinities_16nm_2d_in1068_fp16.onnx", "setup_53_mito_affinities_16nm_2d_in1068_fp32.onnx", "setup_53_mito_affinities_16nm_2d_in300_fp16.onnx", "setup_53_mito_affinities_16nm_2d_in300_fp32.onnx", "setup_53_mito_affinities_16nm_2d_in556_fp16.onnx", "setup_53_mito_affinities_16nm_2d_in556_fp32.onnx" ], "frozen": [ "encoder0.bn1.running_mean", "encoder0.bn1.running_var", "encoder0.bn2.running_mean", "encoder0.bn2.running_var", "encoder1.bn1.running_mean", "encoder1.bn1.running_var", "encoder1.bn2.running_mean", "encoder1.bn2.running_var", "encoder2.bn1.running_mean", "encoder2.bn1.running_var", "encoder2.bn2.running_mean", "encoder2.bn2.running_var", "bottleneck.bn1.running_mean", "bottleneck.bn1.running_var", "bottleneck.bn2.running_mean", "bottleneck.bn2.running_var", "decoder2.bn1.running_mean", "decoder2.bn1.running_var", "decoder2.bn2.running_mean", "decoder2.bn2.running_var", "decoder1.bn1.running_mean", "decoder1.bn1.running_var", "decoder1.bn2.running_mean", "decoder1.bn2.running_var", "decoder0.bn1.running_mean", "decoder0.bn1.running_var", "decoder0.bn2.running_mean", "decoder0.bn2.running_var" ], "note": "params named as the ONNX initializers; BatchNorm running stats frozen", "spike": { "setup": "TF.js 4.22.0 WebGPU backend, Chrome 154 on macOS, Apple GPU (metal-3, unified memory); batch 1, fp32, masked BCE on logits, Adam lr 1e-4; graph = finetune/train_graph.json, weights read from the served fp32 .onnx; median of 8 steps after 2 warm-up; peak = tf.profile peakBytes above baseline (tracked tensors; the WebGPU buffer pool holds more)", "date": "2026-10-17", "params": 1927907, "full": { "300": { "step_ms": 308, "peak_mb": 980 }, "556": { "step_ms": 1189, "peak_mb": 3897 } }, "head_through_full_graph": { "300": { "step_ms": 31 }, "556": { "step_ms": 133 } }, "forward_ms": { "300": 11, "556": 41 }, "tfjs_webgpu_vs_ort_webgpu_556": { "max_abs": 1.81e-05, "flips": 0 }, "note": "2026-10-17 rerun of finetune_spike.html unchanged: peak_mb corrected (the 2026-10-16 values, 320/1307 MB for setup_53_2d, were not reproducible; this rerun matches an independent RTX 2080 Ti run: 980/3897 MB). Peak counts tracked tensors only, the WebGPU buffer pool holds more." } } }, "source_repo": "cellmap/mito-aff-unetiso-setup-53-2d" }