{
  "model": "setup_55",
  "task": "nuc",
  "architecture": "Unet3dIso",
  "checkpoint": "model_checkpoint_330000",
  "params": 5602529,
  "channels": [
    "nuc"
  ],
  "input_voxel_size_nm": [
    128,
    128,
    128
  ],
  "output_voxel_size_nm": [
    128,
    128,
    128
  ],
  "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,z,y,x]; output 'output' float32 [1,C,z',y',x']",
  "variants": [
    {
      "model": "setup_55",
      "variant": "fp32",
      "input_size": 172,
      "output_size": 84,
      "file_mb": 22.51,
      "file": "setup_55_nuc_128nm_in172_fp32.onnx",
      "max_abs": 2.193450927734375e-05,
      "mean_abs": 2.589692257970455e-06,
      "flip_frac": 0.0,
      "ort_cpu_s": 1.19
    },
    {
      "model": "setup_55",
      "variant": "fp16",
      "input_size": 172,
      "output_size": 84,
      "file_mb": 11.32,
      "file": "setup_55_nuc_128nm_in172_fp16.onnx",
      "max_abs": 0.13435745239257812,
      "mean_abs": 0.01906723529100418,
      "flip_frac": 0.0002345184105388187,
      "ort_cpu_s": 2.42
    },
    {
      "model": "setup_55",
      "variant": "int8w",
      "input_size": 172,
      "output_size": 84,
      "file_mb": 5.76,
      "file": "setup_55_nuc_128nm_in172_int8w.onnx",
      "max_abs": 0.5892090797424316,
      "mean_abs": 0.09005698561668396,
      "flip_frac": 0.0009954378576827556,
      "ort_cpu_s": 1.1
    },
    {
      "model": "setup_55",
      "variant": "int8w_fp16",
      "input_size": 172,
      "output_size": 84,
      "file_mb": 5.77,
      "file": "setup_55_nuc_128nm_in172_int8w_fp16.onnx",
      "max_abs": 0.5301542282104492,
      "mean_abs": 0.08458311855792999,
      "flip_frac": 0.0009566326530612245,
      "ort_cpu_s": 2.14
    },
    {
      "model": "setup_55",
      "variant": "fp32",
      "input_size": 140,
      "output_size": 52,
      "file_mb": 22.51,
      "file": "setup_55_nuc_128nm_in140_fp32.onnx",
      "max_abs": 2.193450927734375e-05,
      "mean_abs": 2.494327873137081e-06,
      "flip_frac": 0.0,
      "ort_cpu_s": 0.52
    },
    {
      "model": "setup_55",
      "variant": "fp16",
      "input_size": 140,
      "output_size": 52,
      "file_mb": 11.32,
      "file": "setup_55_nuc_128nm_in140_fp16.onnx",
      "max_abs": 0.08014488220214844,
      "mean_abs": 0.018255988135933876,
      "flip_frac": 3.555985434683659e-05,
      "ort_cpu_s": 1.21
    },
    {
      "model": "setup_55",
      "variant": "int8w",
      "input_size": 140,
      "output_size": 52,
      "file_mb": 5.76,
      "file": "setup_55_nuc_128nm_in140_int8w.onnx",
      "max_abs": 0.4925203323364258,
      "mean_abs": 0.08849068731069565,
      "flip_frac": 0.0002204710969503869,
      "ort_cpu_s": 0.53
    },
    {
      "model": "setup_55",
      "variant": "int8w_fp16",
      "input_size": 140,
      "output_size": 52,
      "file_mb": 5.77,
      "file": "setup_55_nuc_128nm_in140_int8w_fp16.onnx",
      "max_abs": 0.47701168060302734,
      "mean_abs": 0.07975325733423233,
      "flip_frac": 0.00017068730086481566,
      "ort_cpu_s": 1.14
    },
    {
      "model": "setup_55",
      "variant": "fp32",
      "input_size": 124,
      "output_size": 36,
      "file_mb": 22.51,
      "file": "setup_55_nuc_128nm_in124_fp32.onnx",
      "max_abs": 1.52587890625e-05,
      "mean_abs": 2.2992719550529728e-06,
      "flip_frac": 0.0,
      "ort_cpu_s": 0.33
    },
    {
      "model": "setup_55",
      "variant": "fp16",
      "input_size": 124,
      "output_size": 36,
      "file_mb": 11.32,
      "file": "setup_55_nuc_128nm_in124_fp16.onnx",
      "max_abs": 0.06598949432373047,
      "mean_abs": 0.01684611104428768,
      "flip_frac": 2.143347050754458e-05,
      "ort_cpu_s": 0.78
    },
    {
      "model": "setup_55",
      "variant": "int8w",
      "input_size": 124,
      "output_size": 36,
      "file_mb": 5.76,
      "file": "setup_55_nuc_128nm_in124_int8w.onnx",
      "max_abs": 0.4925193786621094,
      "mean_abs": 0.08635841310024261,
      "flip_frac": 0.0,
      "ort_cpu_s": 0.3
    },
    {
      "model": "setup_55",
      "variant": "int8w_fp16",
      "input_size": 124,
      "output_size": 36,
      "file_mb": 5.77,
      "file": "setup_55_nuc_128nm_in124_int8w_fp16.onnx",
      "max_abs": 0.47701072692871094,
      "mean_abs": 0.0763426423072815,
      "flip_frac": 0.0,
      "ort_cpu_s": 0.72
    }
  ],
  "source_repo": "cellmap/nuc-unetiso-setup-55",
  "torch_version": "2.9.1",
  "exporter": "torch.onnx.export (TorchScript, dynamo=False)",
  "opset": 18,
  "finetune": {
    "head": {
      "final_layer": "final_conv.convs.0",
      "onnx_node": "/final_conv/convs.0/Conv",
      "features": "/final_conv/Slice_output_0",
      "weight": "final_conv.convs.0.weight",
      "bias": "final_conv.convs.0.bias",
      "in_channels": 32,
      "out_channels": 1,
      "activation": "none (logits)",
      "files": [
        "setup_55_nuc_128nm_in124_fp16.onnx",
        "setup_55_nuc_128nm_in124_fp32.onnx",
        "setup_55_nuc_128nm_in140_fp16.onnx",
        "setup_55_nuc_128nm_in140_fp32.onnx",
        "setup_55_nuc_128nm_in172_fp16.onnx",
        "setup_55_nuc_128nm_in172_fp32.onnx"
      ],
      "initializers": "finetune/initializers.json",
      "loss": "masked BCE on logits",
      "note": "train the final 1x1 conv on frozen features",
      "torch_weight": {
        "torch": "final_conv.weight",
        "transform": "[:, :, 0]",
        "to_onnx": [
          {
            "index": [
              2,
              0
            ]
          }
        ]
      },
      "torch_bias": {
        "torch": "final_conv.bias",
        "transform": "identity",
        "to_onnx": []
      },
      "weight_map": "finetune/weight_map.json"
    },
    "full": {
      "library": "tfjs",
      "files": [
        "finetune/train_graph.json",
        "finetune/cellmap_tfjs_train.js"
      ],
      "weight_map": "finetune/weight_map.json",
      "torch_params": "finetune/torch_params.json",
      "initializers": "finetune/initializers.json",
      "write_back": [
        "setup_55_nuc_128nm_in124_fp16.onnx",
        "setup_55_nuc_128nm_in124_fp32.onnx",
        "setup_55_nuc_128nm_in140_fp16.onnx",
        "setup_55_nuc_128nm_in140_fp32.onnx",
        "setup_55_nuc_128nm_in172_fp16.onnx",
        "setup_55_nuc_128nm_in172_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": "NDHWC graph, params named as the PyTorch state_dict (Conv3d [O,I,kd,kh,kw], ConvTranspose3d [I,O,2,2,2]); read them from the .onnx with torch_params.json from_onnx, write back with weight_map.json to_onnx; BatchNorm running stats frozen. Plain TF.js autodiff keeps every intermediate until the backward pass: use the smallest block (3D memory grows with the block volume)"
    }
  }
}