finetune runtime: note on Blob import (HF serves .js as non-JS MIME)

fb68001 verified 12 days ago

4.01 kB

<!doctype html>
<meta charset="utf-8">
<title>cellmap-flow finetune spike</title>
<!-- Times one full-network training step (masked BCE on logits, Adam, batch 1, fp32) of a 2D student in TF.js WebGPU,
     built from the model folder's finetune/ files, and checks the TF.js forward against onnxruntime-web.
     Open over https (WebGPU needs a secure context), keep the tab in the foreground.
     ?model=setup_53_2d&sizes=300,556  (repo cellmap/onnx-web-models) -->
<pre id="log">loading…</pre>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@4.22.0/dist/tf.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-webgpu@4.22.0/dist/tf-backend-webgpu.min.js"></script>
<script type="module">
const q = new URLSearchParams(location.search);
const model = q.get("model") || "setup_53_2d", sizes = (q.get("sizes") || "300,556").split(",").map(Number);
const repo = q.get("repo") || "https://huggingface.co/cellmap/onnx-web-models/resolve/main/";
const base = repo + model + "/";
const log = s => { document.getElementById("log").textContent += "\n" + s; console.log(s); };
// HF serves .js without a JavaScript MIME type, so import it from a Blob
const rt = await import(URL.createObjectURL(new Blob([await (await fetch(base + "finetune/cellmap_tfjs_train.js")).text()],
                                                     { type: "text/javascript" })));
await tf.setBackend("webgpu"); await tf.ready();
const man = await (await fetch(base + "manifest.json")).json();
const graph = await (await fetch(base + "finetune/train_graph.json")).json();
const offsAll = await (await fetch(base + "finetune/initializers.json")).json();
const ort = await import("https://cdn.jsdelivr.net/npm/onnxruntime-web@1.27.0-dev.20260506-673c3320fc/dist/ort.webgpu.bundle.min.mjs");
const results = [];
for (const N of sizes) {
  const file = man.variants.find(v => v.variant === "fp32" && v.input_size[1] === N)?.file;
  if (!file) { log(`no fp32 file for ${N}`); continue; }
  const bytes = new Uint8Array(await (await fetch(base + file)).arrayBuffer());
  const net = rt.buildNet(tf, graph, rt.readOnnxWeights(bytes, offsAll[file], Object.keys(graph.params)));
  const M = N - graph.context;
  // forward check vs onnxruntime-web on a smooth synthetic slice
  const xa = new Float32Array(N * N).map((_, i) => Math.sin(Math.floor(i / N) / 9) * Math.cos((i % N) / 13) * 0.8);
  const s = await ort.InferenceSession.create(bytes, { executionProviders: ["webgpu"] });
  const ro = (await s.run({ input: new ort.Tensor("float32", xa, [1, 1, 1, N, N]) })).output.data;
  const zd = await tf.tidy(() => tf.transpose(net.forward(tf.tensor(xa, [1, N, N, 1])), [0, 3, 1, 2])).data();
  let maxAbs = 0, flips = 0;
  zd.forEach((v, i) => { maxAbs = Math.max(maxAbs, Math.abs(v - ro[i])); flips += (v > 0) !== (ro[i] > 0); });
  // training steps
  const C = man.channels.length;
  const x = tf.randomUniform([1, N, N, 1], -1, 1);
  const t = tf.cast(tf.greater(tf.randomUniform([1, M, M, C]), 0.5), "float32");
  const m = tf.cast(tf.greater(tf.randomUniform([1, M, M, C]), 0.7), "float32");
  const opt = tf.train.adam(1e-4);
  const step = () => rt.trainStep(tf, net, opt, x, t, m);
  for (let i = 0; i < 2; i++) { const l = step(); await l.data(); l.dispose(); }
  const base0 = tf.memory().numBytes;
  const prof = await tf.profile(async () => { const l = step(); await l.data(); l.dispose(); });
  const times = [];
  for (let i = 0; i < 8; i++) { const t0 = performance.now(); const l = step(); await l.data(); times.push(performance.now() - t0); l.dispose(); }
  times.sort((a, b) => a - b);
  const r = { model, N, step_ms_median: Math.round(times[4]), step_ms_min: Math.round(times[0]),
              train_peak_mb: Math.round((prof.peakBytes - base0) / 1e6), vs_ort_max_abs: maxAbs, vs_ort_flips: flips };
  results.push(r); log(JSON.stringify(r));
  [x, t, m].forEach(a => a.dispose()); opt.dispose(); net.vars.forEach(v => v.dispose());
}
window.SPIKE = results;
log("done");
</script>

	<!doctype html>
	<meta charset="utf-8">
	<title>cellmap-flow finetune spike</title>
	<!-- Times one full-network training step (masked BCE on logits, Adam, batch 1, fp32) of a 2D student in TF.js WebGPU,
	built from the model folder's finetune/ files, and checks the TF.js forward against onnxruntime-web.
	Open over https (WebGPU needs a secure context), keep the tab in the foreground.
	?model=setup_53_2d&sizes=300,556 (repo cellmap/onnx-web-models) -->
	<pre id="log">loading…</pre>
	<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@4.22.0/dist/tf.min.js"></script>
	<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-webgpu@4.22.0/dist/tf-backend-webgpu.min.js"></script>
	<script type="module">
	const q = new URLSearchParams(location.search);
	const model = q.get("model") \|\| "setup_53_2d", sizes = (q.get("sizes") \|\| "300,556").split(",").map(Number);
	const repo = q.get("repo") \|\| "https://huggingface.co/cellmap/onnx-web-models/resolve/main/";
	const base = repo + model + "/";
	const log = s => { document.getElementById("log").textContent += "\n" + s; console.log(s); };
	// HF serves .js without a JavaScript MIME type, so import it from a Blob
	const rt = await import(URL.createObjectURL(new Blob([await (await fetch(base + "finetune/cellmap_tfjs_train.js")).text()],
	{ type: "text/javascript" })));
	await tf.setBackend("webgpu"); await tf.ready();
	const man = await (await fetch(base + "manifest.json")).json();
	const graph = await (await fetch(base + "finetune/train_graph.json")).json();
	const offsAll = await (await fetch(base + "finetune/initializers.json")).json();
	const ort = await import("https://cdn.jsdelivr.net/npm/onnxruntime-web@1.27.0-dev.20260506-673c3320fc/dist/ort.webgpu.bundle.min.mjs");
	const results = [];
	for (const N of sizes) {
	const file = man.variants.find(v => v.variant === "fp32" && v.input_size[1] === N)?.file;
	if (!file) { log(`no fp32 file for ${N}`); continue; }
	const bytes = new Uint8Array(await (await fetch(base + file)).arrayBuffer());
	const net = rt.buildNet(tf, graph, rt.readOnnxWeights(bytes, offsAll[file], Object.keys(graph.params)));
	const M = N - graph.context;
	// forward check vs onnxruntime-web on a smooth synthetic slice
	const xa = new Float32Array(N * N).map((_, i) => Math.sin(Math.floor(i / N) / 9) * Math.cos((i % N) / 13) * 0.8);
	const s = await ort.InferenceSession.create(bytes, { executionProviders: ["webgpu"] });
	const ro = (await s.run({ input: new ort.Tensor("float32", xa, [1, 1, 1, N, N]) })).output.data;
	const zd = await tf.tidy(() => tf.transpose(net.forward(tf.tensor(xa, [1, N, N, 1])), [0, 3, 1, 2])).data();
	let maxAbs = 0, flips = 0;
	zd.forEach((v, i) => { maxAbs = Math.max(maxAbs, Math.abs(v - ro[i])); flips += (v > 0) !== (ro[i] > 0); });
	// training steps
	const C = man.channels.length;
	const x = tf.randomUniform([1, N, N, 1], -1, 1);
	const t = tf.cast(tf.greater(tf.randomUniform([1, M, M, C]), 0.5), "float32");
	const m = tf.cast(tf.greater(tf.randomUniform([1, M, M, C]), 0.7), "float32");
	const opt = tf.train.adam(1e-4);
	const step = () => rt.trainStep(tf, net, opt, x, t, m);
	for (let i = 0; i < 2; i++) { const l = step(); await l.data(); l.dispose(); }
	const base0 = tf.memory().numBytes;
	const prof = await tf.profile(async () => { const l = step(); await l.data(); l.dispose(); });
	const times = [];
	for (let i = 0; i < 8; i++) { const t0 = performance.now(); const l = step(); await l.data(); times.push(performance.now() - t0); l.dispose(); }
	times.sort((a, b) => a - b);
	const r = { model, N, step_ms_median: Math.round(times[4]), step_ms_min: Math.round(times[0]),
	train_peak_mb: Math.round((prof.peakBytes - base0) / 1e6), vs_ort_max_abs: maxAbs, vs_ort_flips: flips };
	results.push(r); log(JSON.stringify(r));
	[x, t, m].forEach(a => a.dispose()); opt.dispose(); net.vars.forEach(v => v.dispose());
	}
	window.SPIKE = results;
	log("done");
	</script>