setup_53_2d/finetune/cellmap_tfjs_train.js

// Reference runtime for cellmap-flow browser finetuning ("full" mode, library tfjs).
// Builds a trainable TF.js network from finetune/train_graph.json, takes its weights from the served .onnx
// (finetune/initializers.json gives the byte offset of every initializer), and writes trained weights back into
// the .onnx bytes in place, so the result is served by onnxruntime-web exactly as before.
//
//   (served from Hugging Face this file has no JavaScript MIME type: fetch it and import() a Blob URL, or bundle it)
//   const g = await (await fetch(base + "finetune/train_graph.json")).json();
//   const offs = (await (await fetch(base + "finetune/initializers.json")).json())[onnxFile];
//   const onnxBytes = new Uint8Array(await (await fetch(base + onnxFile)).arrayBuffer());
//   const net = buildNet(tf, g, readOnnxWeights(onnxBytes, offs, Object.keys(g.params)));
//   const opt = tf.train.adam(1e-4);
//   const loss = trainStep(tf, net, opt, x, target, mask);   // x [1,H,W,1]; target/mask [1,H-ctx,W-ctx,C]
//   const newBytes = writeOnnxWeights(onnxBytes, offs, net.exportWeights());
//
// Tensors are NHWC. ONNX conv weights [O,I,kh,kw] -> tf [kh,kw,I,O]; ConvTranspose [I,O,kh,kw] -> tf [kh,kw,O,I].
//
// 3D models (train_graph layout "NDHWC"): params are the PyTorch state_dict tensors, rebuilt from the served .onnx
// (Conv3d split per depth tap, ConvTranspose3d as 1x1 conv) with finetune/torch_params.json, and written back with
// finetune/weight_map.json:
//   const tp = await (await fetch(base + "finetune/torch_params.json")).json();
//   const wm = await (await fetch(base + "finetune/weight_map.json")).json();
//   const net = buildNet(tf, g, readTorchParams(onnxBytes, offs, tp));          // x [1,D,H,W,1]
//   const newBytes = writeOnnxWeights(onnxBytes, offs, toOnnxWeights(wm, net.exportWeights(), offs));
// WebGPU workarounds in the 3D path: bias / norm affine on rank-2 views (a rank-5 add with a short 1-D operand fails
// to compile), crops as rank-4/3 slices (rank-5 slice gradient fails), ConvTranspose3d k=s=2 as matmul +
// depth-to-space (conv3dTranspose has no gradient), nearest upsample as rank-3 tiles. Plain autodiff keeps every
// intermediate: only small blocks fit (see manifest note).

function f16ToF32(h) {
  const s = (h & 0x8000) ? -1 : 1, e = (h >> 10) & 0x1f, f = h & 0x3ff;
  if (e === 0) return s * Math.pow(2, -14) * (f / 1024);
  if (e === 31) return f ? NaN : s * Infinity;
  return s * Math.pow(2, e - 15) * (1 + f / 1024);
}
const _f = new Float32Array(1), _u = new Uint32Array(_f.buffer);
function f32ToF16(v) {
  _f[0] = v; const x = _u[0];
  const s = (x >>> 16) & 0x8000; let e = ((x >>> 23) & 0xff) - 112; let m = x & 0x7fffff;
  if (e <= 0) { if (e < -10) return s; m |= 0x800000; const t = 14 - e; const r = (m >> t) + ((m >> (t - 1)) & 1); return s | r; }
  if (e >= 31) return s | 0x7c00 | (((x >>> 23) & 0xff) === 0xff && m ? 0x200 : 0);
  let r = s | (e << 10) | (m >> 13);
  if (m & 0x1000) r += 1;  // round half up on the dropped bits (good enough for weights)
  return r;
}

export function readOnnxWeights(bytes, offs, names) {
  const out = {};
  for (const n of names) {
    const o = offs[n];
    if (!o) throw new Error(`initializer ${n} not in file`);
    const dv = new DataView(bytes.buffer, bytes.byteOffset + o.offset, o.nbytes);
    if (o.dtype === "float32") {
      const a = new Float32Array(o.nbytes / 4);
      for (let i = 0; i < a.length; i++) a[i] = dv.getFloat32(4 * i, true);
      out[n] = a;
    } else if (o.dtype === "float16") {
      const a = new Float32Array(o.nbytes / 2);
      for (let i = 0; i < a.length; i++) a[i] = f16ToF32(dv.getUint16(2 * i, true));
      out[n] = a;
    } else throw new Error(`${n}: dtype ${o.dtype} is not writable`);
  }
  return out;
}

export function writeOnnxWeights(bytes, offs, weights) {
  const b = new Uint8Array(bytes);  // copy
  for (const [n, a] of Object.entries(weights)) {
    const o = offs[n];
    const dv = new DataView(b.buffer, b.byteOffset + o.offset, o.nbytes);
    if (o.dtype === "float32") { if (a.length * 4 !== o.nbytes) throw new Error(n); a.forEach((v, i) => dv.setFloat32(4 * i, v, true)); }
    else if (o.dtype === "float16") { if (a.length * 2 !== o.nbytes) throw new Error(n); a.forEach((v, i) => dv.setUint16(2 * i, f32ToF16(v), true)); }
    else throw new Error(`${n}: dtype ${o.dtype} is not writable`);
  }
  return b;
}

export function buildNet(tf, graph, weights, { trainable = null } = {}) {
  // trainable: optional Set of param names to train (default: graph's trainable flags); e.g. head only.
  if (graph.layout === "NDHWC") return buildNet3d(tf, graph, weights, { trainable });
  const P = {};
  const kind = {};
  for (const op of graph.ops) {
    if (op.op === "conv") kind[op.w] = "conv";
    if (op.op === "conv_t") kind[op.w] = "conv_t";
  }
  for (const [n, meta] of Object.entries(graph.params)) {
    let t = tf.tensor(weights[n], meta.shape, "float32");
    if (kind[n] === "conv") t = tf.transpose(t, [2, 3, 1, 0]);
    else if (kind[n] === "conv_t") t = tf.transpose(t, [2, 3, 1, 0]);
    const train = trainable ? trainable.has(n) : meta.trainable;
    P[n] = train ? tf.variable(t, true) : t;  // unnamed: tf variable names are global, several nets may coexist
    if (train) t.dispose();
  }
  function forward(x) {
    const T = { [graph.input]: x };
    for (const op of graph.ops) {
      const a = T[op.x];
      let y;
      switch (op.op) {
        case "conv": y = tf.conv2d(a, P[op.w], op.stride, "valid"); if (op.b) y = tf.add(y, P[op.b]); break;
        case "conv_t": {
          const [kh, kw, o] = P[op.w].shape, [n, h, w] = a.shape;
          y = tf.conv2dTranspose(a, P[op.w], [n, (h - 1) * op.stride[0] + kh, (w - 1) * op.stride[1] + kw, o], op.stride, "valid");
          if (op.b) y = tf.add(y, P[op.b]); break;
        }
        case "maxpool": y = tf.maxPool(a, op.k, op.stride, "valid"); break;
        case "bn": y = tf.batchNorm(a, P[op.mean], P[op.var], P[op.bias], P[op.scale], op.eps); break;  // frozen stats
        case "inorm": { const { mean, variance } = tf.moments(a, [1, 2], true);
          y = tf.add(tf.mul(tf.div(tf.sub(a, mean), tf.sqrt(tf.add(variance, op.eps))), P[op.scale]), P[op.bias]); break; }
        case "relu": y = tf.relu(a); break;
        case "add": y = tf.add(a, T[op.x2]); break;
        case "concat": y = tf.concat(op.xs.map(k => T[k]), 3); break;
        case "crop": { const s = a.shape, ax = op.axis === "h" ? 1 : 2, begin = [0, 0, 0, 0], size = [-1, -1, -1, -1];
          begin[ax] = op.start; size[ax] = s[ax] - op.start - op.end_margin; y = tf.slice(a, begin, size); break; }
        case "upsample": { const [n, h, w, c] = a.shape, f = op.factor;  // nearest = repeat; rank-3 tiles (tf has no
          const yw = tf.reshape(tf.tile(tf.reshape(a, [n * h * w, 1, c]), [1, f, 1]), [n * h, 1, w * f * c]);  // rank-6 tile grad)
          y = tf.reshape(tf.tile(yw, [1, f, 1]), [n, h * f, w * f, c]); break; }
        default: throw new Error(`op ${op.op}`);
      }
      T[op.y] = y;
    }
    return T[graph.output];
  }
  function exportWeights() {  // back to ONNX layout, trainable params only
    const out = {};
    for (const [n, t] of Object.entries(P)) {
      if (!(t instanceof tf.Variable)) continue;
      const v = kind[n] ? tf.transpose(t, [3, 2, 0, 1]) : t;
      out[n] = v.dataSync().slice();
      if (v !== t) v.dispose();
    }
    return out;
  }
  const vars = Object.values(P).filter(t => t instanceof tf.Variable);
  return { forward, exportWeights, vars, params: P, context: graph.context };
}

// masked BCE on logits: mean over mask>0 of max(z,0) - z*t + log(1+exp(-|z|))
export function maskedBce(tf, z, t, m) {
  return tf.tidy(() => {
    const l = tf.add(tf.sub(tf.relu(z), tf.mul(z, t)), tf.log1p(tf.exp(tf.neg(tf.abs(z)))));
    return tf.div(tf.sum(tf.mul(l, m)), tf.maximum(tf.sum(m), 1));
  });
}

export function trainStep(tf, net, opt, x, target, mask) {
  return opt.minimize(() => maskedBce(tf, net.forward(x), target, mask), true, net.vars);
}

// ---------------------------------------------------------------- structured weight transforms (3D folders)
// ops on {data: Float32Array, shape}: index [axis,k] | permute [...] | reshape [...] | repeat_interleave [axis,n];
// recipes may start with {from: name} or {stack: [names], axis}. Same semantics as finetune_spec.py apply_ops.
const prod = a => a.reduce((x, y) => x * y, 1);
function nd(a, ops, get) {
  for (const o of ops) {
    if ("from" in o) a = get(o.from);
    else if ("stack" in o) {
      const parts = o.stack.map(get), ax = o.axis, s = parts[0].shape, outer = prod(s.slice(0, ax)), inner = prod(s.slice(ax));
      const d = new Float32Array(outer * inner * parts.length);
      for (let i = 0; i < outer; i++) parts.forEach((p, k) => d.set(p.data.subarray(i * inner, (i + 1) * inner), (i * parts.length + k) * inner));
      a = { data: d, shape: [...s.slice(0, ax), parts.length, ...s.slice(ax)] };
    } else if ("index" in o) {
      const [ax, k] = o.index, s = a.shape, outer = prod(s.slice(0, ax)), inner = prod(s.slice(ax + 1)), n = s[ax];
      const d = new Float32Array(outer * inner);
      for (let i = 0; i < outer; i++) d.set(a.data.subarray((i * n + k) * inner, (i * n + k + 1) * inner), i * inner);
      a = { data: d, shape: [...s.slice(0, ax), ...s.slice(ax + 1)] };
    } else if ("reshape" in o) {
      if (prod(o.reshape) !== a.data.length) throw new Error("reshape size");
      a = { data: a.data, shape: o.reshape.slice() };
    } else if ("permute" in o) {
      const p = o.permute, s = a.shape, r = s.length, os = p.map(i => s[i]);
      const st = new Array(r); st[r - 1] = 1; for (let i = r - 2; i >= 0; i--) st[i] = st[i + 1] * s[i + 1];
      const pst = p.map(i => st[i]), d = new Float32Array(a.data.length), idx = new Array(r).fill(0);
      for (let j = 0; j < d.length; j++) {
        let src = 0; for (let i = 0; i < r; i++) src += idx[i] * pst[i];
        d[j] = a.data[src];
        for (let i = r - 1; i >= 0; i--) { if (++idx[i] < os[i]) break; idx[i] = 0; }
      }
      a = { data: d, shape: os };
    } else if ("repeat_interleave" in o) {
      const [ax, n] = o.repeat_interleave, s = a.shape, outer = prod(s.slice(0, ax + 1)), inner = prod(s.slice(ax + 1));
      const d = new Float32Array(a.data.length * n);
      for (let i = 0; i < outer; i++) for (let k = 0; k < n; k++) d.set(a.data.subarray(i * inner, (i + 1) * inner), (i * n + k) * inner);
      a = { data: d, shape: s.map((v, i) => (i === ax ? v * n : v)) };
    } else throw new Error(`transform ${JSON.stringify(o)}`);
  }
  return a;
}

export function readTorchParams(bytes, offs, torchParams) {
  // {torch_name: Float32Array (PyTorch layout)} from one served .onnx, via torch_params.json from_onnx recipes
  const names = new Set();
  for (const p of Object.values(torchParams)) for (const o of p.from_onnx) (o.from ? [o.from] : o.stack || []).forEach(n => names.add(n));
  const raw = readOnnxWeights(bytes, offs, [...names]);
  const get = n => ({ data: raw[n], shape: offs[n].shape });
  const out = {};
  for (const [k, p] of Object.entries(torchParams)) {
    const a = nd(null, p.from_onnx, get);
    if (a.data.length !== prod(p.shape)) throw new Error(k);
    out[k] = a.data;
  }
  return out;
}

export function toOnnxWeights(weightMap, params, offs) {
  // trained PyTorch-layout params (net.exportWeights()) -> {onnx_initializer: Float32Array} for writeOnnxWeights,
  // only initializers of trained params that are float in this file (offs)
  const out = {};
  for (const [n, e] of Object.entries(weightMap)) {
    const o = offs[n];
    if (!(e.torch in params) || !o || (o.dtype !== "float32" && o.dtype !== "float16")) continue;
    const p = params[e.torch];
    out[n] = nd(p, e.to_onnx).data;
  }
  return out;
}

function buildNet3d(tf, graph, weights, { trainable = null } = {}) {
  const P = {}, kind = {}, shape = {};
  for (const op of graph.ops) {
    if (op.op === "conv") kind[op.w] = "conv";
    if (op.op === "conv_t") { kind[op.w] = "conv_t"; kind[op.b] = "conv_t_b"; }
  }
  for (const [n, meta] of Object.entries(graph.params)) {
    shape[n] = meta.shape;
    let t = tf.tensor(weights[n], meta.shape, "float32");
    if (kind[n] === "conv") t = tf.transpose(t, [2, 3, 4, 1, 0]);  // [O,I,kd,kh,kw] -> [kd,kh,kw,I,O]
    else if (kind[n] === "conv_t") t = tf.reshape(t, [meta.shape[0], meta.shape[1] * 8]);  // [I,O,2,2,2] -> [I, O*8]
    const train = trainable ? trainable.has(n) : meta.trainable;
    P[n] = train ? tf.variable(t, true) : t;
    if (train) t.dispose();
  }
  const flat = (a, f) => { const s = a.shape, c = s[s.length - 1]; return tf.reshape(f(tf.reshape(a, [-1, c])), s); };
  const addBias = (a, b) => flat(a, v => tf.add(v, b));
  function crop(a, ax, start, endMargin) {  // ax 1|2|3 (d|h|w) of [N,D,H,W,C] via a rank-4 / rank-3 view
    const [n, d, h, w, c] = a.shape, len = a.shape[ax] - start - endMargin;
    const v = ax === 1 ? [n, d, h * w, c] : ax === 2 ? [n * d, h, w, c] : [n * d * h, w, c];
    const begin = v.map(() => 0), size = v.map(() => -1);
    begin[1] = start; size[1] = len;
    const out = a.shape.slice(); out[ax] = len;
    return tf.reshape(tf.slice(tf.reshape(a, v), begin, size), out);
  }
  function repeatAxis(a, ax, f) {  // nearest upsample along one axis: rank-3 tile
    const s = a.shape, outer = prod(s.slice(0, ax + 1)), inner = prod(s.slice(ax + 1));
    const out = s.slice(); out[ax] *= f;
    return tf.reshape(tf.tile(tf.reshape(a, [outer, 1, inner]), [1, f, 1]), out);
  }
  function forward(x) {
    const T = { [graph.input]: x };
    for (const op of graph.ops) {
      const a = T[op.x];
      let y;
      switch (op.op) {
        case "conv": y = tf.conv3d(a, P[op.w], op.stride, "valid"); if (op.b) y = addBias(y, P[op.b]); break;
        case "conv_t": {  // k = s = 2: [P,I] x [I, O*8] (+bias repeated 8x), then depth-to-space (rank <= 6)
          const [n, d, h, w, ci] = a.shape, o = P[op.w].shape[1] / 8;
          if (op.stride.some(s => s !== 2) || shape[op.w].slice(2).some(k => k !== 2)) throw new Error("conv_t: k=s=2 only");
          let z = tf.matMul(tf.reshape(a, [-1, ci]), P[op.w]);
          if (op.b) z = tf.add(z, tf.reshape(tf.tile(tf.reshape(P[op.b], [o, 1]), [1, 8]), [o * 8]));
          z = tf.transpose(tf.reshape(z, [n * d, h, w, o, 2, 4]), [0, 4, 1, 2, 5, 3]);  // [nd,a,h,w,(b e),o]
          z = tf.transpose(tf.reshape(z, [n * d * 2 * h, w, 2, 2, o]), [0, 2, 1, 3, 4]);  // [nd a h, b, w, e, o]
          y = tf.reshape(z, [n, 2 * d, 2 * h, 2 * w, o]); break;
        }
        case "maxpool": y = tf.maxPool3d(a, op.k, op.stride, "valid"); break;
        case "bn": y = flat(a, v => tf.batchNorm(v, P[op.mean], P[op.var], P[op.bias], P[op.scale], op.eps)); break;
        case "inorm": { const { mean, variance } = tf.moments(a, [1, 2, 3], true);
          y = flat(tf.div(tf.sub(a, mean), tf.sqrt(tf.add(variance, op.eps))), v => tf.add(tf.mul(v, P[op.scale]), P[op.bias])); break; }
        case "relu": y = tf.relu(a); break;
        case "add": y = tf.add(a, T[op.x2]); break;
        case "concat": y = tf.concat(op.xs.map(k => T[k]), 4); break;
        case "crop": y = crop(a, 1 + "dhw".indexOf(op.axis), op.start, op.end_margin); break;
        case "upsample": y = [1, 2, 3].reduce((v, ax) => repeatAxis(v, ax, op.factor), a); break;
        default: throw new Error(`op ${op.op}`);
      }
      T[op.y] = y;
    }
    return T[graph.output];
  }
  function exportWeights() {  // trainable params, PyTorch layout, keyed by state_dict name
    const out = {};
    for (const [n, t] of Object.entries(P)) {
      if (!(t instanceof tf.Variable)) continue;
      const v = kind[n] === "conv" ? tf.transpose(t, [4, 3, 0, 1, 2]) : t;  // conv_t: [I, O*8] is already [I,O,2,2,2] flat
      out[n] = { data: v.dataSync().slice(), shape: shape[n] };
      if (v !== t) v.dispose();
    }
    return out;
  }
  const vars = Object.values(P).filter(t => t instanceof tf.Variable);
  return { forward, exportWeights, vars, params: P, context: graph.context };
}