diff --git a/CHANGELOG.md b/CHANGELOG.md index 5980c83e..0ed67bc6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ All notable changes to cellmap-flow. The format follows [Keep a Changelog](https - **Loaded models, at the top of the Models tab, each with a ✕.** The server dashboard and the browser app list every model running or starting, wherever it was picked, and the ✕ beside one stops that model alone, leaving the rest of the tab as it is (`GET /api/models/loaded`, `POST /api/models/remove`). A finetune job's model keeps its ✕ disabled while the job, its server, is alive. See *Usage* in the README. - **Finetuning in the browser.** The app's Finetune tab paints and trains in the tab, with nothing to install: *New Volume* adds a paint layer kept in the browser, and *Start Finetuning* fits the model's last layer to the paint in seconds, with the server's losses, distillation and *Mark as Good*, then serves the result beside the base model, keeps it across visits and downloads it as an `.onnx` (the paint as a zarr). AI-assisted annotation works in the tab too, the page calling Google with the key or token you paste. See *Finetuning in the browser* in the docs' *In the browser* page, and `web/README.md` for how it is built. - **Whole-network finetuning in the browser, for the 2D models.** The Finetune tab offers what each model's manifest on Hugging Face declares, so for `setup_53_2d` and `setup_54_2d` a *Training* choice adds *Whole network*: every weight trains on the paint with TensorFlow.js on the GPU (WebGPU), on Marwan Zouinkhi's training graphs and runtime, and the trained weights are written back into the served `.onnx`. The manifests' `source_repo`, `distilled_from_repo`, `torch_version`, `exporter` and `opset` are kept and shown in the Models tab's search, and the catalog's cached listing is refreshed so they show. See *Training the whole network* in the docs' *In the browser* page, and `web/ARCHITECTURE.md`. +- **Whole-network finetuning in the browser takes the last layer's loss settings, and augmentation.** The loss (BCE, MSE or margin), its margin, label smoothing, class balance and distillation (λ and scope) now apply to the whole network too, where it trained with masked BCE alone; distillation holds the logits to the starting network's on the same tile. A new *Augmentation* box, off by default, gives each step's tiles the server's brightness, noise, flips and quarter turns, with the targets made from the turned paint, so affinities keep their meaning. The run's log says which settings it trained with, and how many tiles were flipped and turned. The last layer is unchanged and takes no augmentation. - **Share carries a model finetuned in the browser, and says what it leaves behind.** A last-layer finetune goes in the link as its base, the file trained on and that layer's weights, and is rebuilt and kept where the link opens; a model whose whole network was trained stays out, as does paint, and a notice under the header says what the link leaves behind and what the opening page left out. Links made before read as they did. See "A link that opens a volume and runs models" in `web/README.md`. - **Label it in the browser.** The Finetune tab's *Patch on screen* tools work in the app as on the server: *Seed from Prediction* copies a model's prediction into the unpainted voxels of the patch on screen (an id per object, background 1), *All Background* labels them background, *Split Objects* gives each connected object an id of its own, and *Undo* takes back the last of them or of an accepted AI annotation (one stack of 10 per volume, where the AI panel had an Undo of its own). *Prediction from* lists the models served in the page, finetuned ones included (default: the volume model's latest finetune), and the seed reads the model's own output (its normalizers, none of its postprocessors) from the chunks already computed for the viewer, on the grid the server would read it on, with the method, threshold, smallest object, connectivity, per-slice and box settings of the server's form. Its methods are *Threshold + components* and, for affinity models, *Mutex watershed*, the browser's segmenters, checked against the server's on fixtures the Python writes; *Distance watershed* and *Model's instances* are not offered, as no browser model outputs a distance or instance ids. - **cellmap-flow in the browser, at https://janelia-cellmap.github.io/cellmap-flow/app/.** The dashboard as a web page, with the models running in the browser: the raw volume is read from its URL (or, in Chrome and Edge, from a folder on the machine, nothing uploaded), ONNX exports of cellmap's networks from Hugging Face ([cellmap/onnx-web-models](https://huggingface.co/cellmap/onnx-web-models)) run with onnxruntime-web on the GPU through WebGPU (on the CPU without it), and the predictions reach Neuroglancer through a service worker, chunk by chunk, as the viewer asks. The Input normalizers and Postprocess steps are the server dashboard's, ported to TypeScript and checked against the Python classes; the on-demand engine is [chunkmirage](https://github.com/JaneliaSciComp/chunkmirage)'s browser engine. Nothing to install and no cluster account, so it is the first place to try cellmap-flow and to test a change: every push to `main` deploys it beside the docs. It works on a phone, where the viewer fills the screen and *Dashboard* opens the controls over it. A link can name the volume and the models to run (`?dataset=&model=setup_53`; the README's example opens an OpenOrganelle volume with `setup_53` running), and *Share* in the header copies such a link to what is open, view and settings included; the address bar itself stays bare, so a new tab starts fresh. A volume on a private network (an institution's `10.x` host) opens as it does in Neuroglancer; if Chrome asks whether the site may use the local network, allow it, and until then the connect panel's error says so. See `web/README.md` and the docs' *In the browser* page. A port of Marwan Zouinkhi's serverless prototype. diff --git a/cellmap_flow/dashboard/static/js/dashboard/finetune/job-card.js b/cellmap_flow/dashboard/static/js/dashboard/finetune/job-card.js index 16cb3934..0d5bb79e 100644 --- a/cellmap_flow/dashboard/static/js/dashboard/finetune/job-card.js +++ b/cellmap_flow/dashboard/static/js/dashboard/finetune/job-card.js @@ -147,6 +147,7 @@ export function createJobCard(lossPlot) { } if (params.loss_type === "margin" && params.margin != null) parts.push(`margin ${params.margin}`); if (params.distillation_lambda != null) parts.push(`distillation ${params.distillation_lambda}`); + if (params.augment === true) parts.push("augmented"); const line = $("jobTraining"); line.textContent = parts.join(", "); line.title = interval diff --git a/cellmap_flow/dashboard/static/js/dashboard/finetune/training-form.js b/cellmap_flow/dashboard/static/js/dashboard/finetune/training-form.js index 8801d72d..bab4d980 100644 --- a/cellmap_flow/dashboard/static/js/dashboard/finetune/training-form.js +++ b/cellmap_flow/dashboard/static/js/dashboard/finetune/training-form.js @@ -135,11 +135,11 @@ function trainingParams() { } // Browser mode: what trains (the Training field); for the head, its // batch in voxels per step, its own field; for the whole network, its - // learning rate, its tile side and the batch above, in tiles. The - // server-only fields above go too and the page's routes ignore them - // (web/src/backend/routes/finetune.ts readTraining). The Loss Function - // offers the head's losses there (BCE, MSE, margin) and is read as on - // the server; the whole network ignores it. + // learning rate, its tile side and the batch above, in tiles, and + // augmentation. The server-only fields above go too and the page's + // routes ignore them (web/src/backend/routes/finetune.ts readTraining). + // The Loss Function offers the head's losses there (BCE, MSE, margin) and + // is read as on the server, by both. if (isBrowserMode()) { const mode = document.getElementById("finetuneMode").value; params.finetune_mode = mode; @@ -192,6 +192,16 @@ export function applyModelDefaults(name) { function updateAugmentAdvice() { const advice = document.getElementById("augmentAdvice"); const epochs = parseInt(document.getElementById("numEpochs").value, 10); + // The browser's whole network trains on the same tiles every epoch, as + // many as the paint makes, which the form cannot know: the advice is the + // epochs alone. + if (isBrowserMode()) { + advice.textContent = epochs > 0 + ? `Each tile is trained on ${epochs} time${epochs === 1 ? "" : "s"}, once an epoch: the more epochs, the more it keeps the network from learning the tiles by heart.` + : ""; + advice.className = "d-block text-muted"; + return; + } const batch = parseInt(document.getElementById("batchSize").value, 10); const patchesRaw = document.getElementById("patchesPerEpoch").value; const patches = patchesRaw === "" || patchesRaw == null diff --git a/cellmap_flow/dashboard/templates/_finetune_tab.html b/cellmap_flow/dashboard/templates/_finetune_tab.html index 1a5f27f0..4efcc22c 100644 --- a/cellmap_flow/dashboard/templates/_finetune_tab.html +++ b/cellmap_flow/dashboard/templates/_finetune_tab.html @@ -541,9 +541,10 @@
Training Configuration
separate fields and one's saved value never lands in the other). A 2D model whose manifest offers it can train its whole network instead (the Training field, shown when the model offers both): that - takes the epochs, its own learning rates, the server's Batch Size (in - tiles) and a Tile Size, and none of the head's loss, distillation, - smoothing or balance fields (head-only), which it does not use. + takes the epochs, the same loss, distillation, smoothing and balance + fields, its own learning rates, the server's Batch Size (in tiles), a + Tile Size and the server's Augmentation (full-also), but not the + head's Voxels per Step (head-only). web/src/browser.css gives the fields the full row. #} {% if mode == "browser" %} @@ -713,8 +714,7 @@
Training Configuration
- {# The whole network trains with masked BCE on the paint alone: the rest are the head's. #} -
+
@@ -766,7 +766,7 @@
Training Configuration
{% if mode == "browser" %} -
-
+
Weights fg and bg equally.
-
+ {# The browser's whole network takes it too (full-also): its tiles are the same every epoch, so it is the case augmentation is for. #} +
+ {% if mode == "browser" %} + + {% else %} + {% endif %}
diff --git a/docs/source/browser.rst b/docs/source/browser.rst index 21799b65..87c884c7 100644 --- a/docs/source/browser.rst +++ b/docs/source/browser.rst @@ -183,12 +183,23 @@ written): serve its fp32 or fp16 file to finetune it. The whole network trains on square tiles cut from the painted slices, in the plane you painted in (x–y, or z–y and z–x where the raw data's voxels are -cubes, as the server's 2D training does), with masked BCE on the painted -voxels and Adam. Its controls are the epochs, its own learning rates (1e-4 by -default, a hundredth of the last layer's), the server's *Batch Size*, in tiles -per step (1 by default), and *Tile Size*. The loss, distillation, smoothing -and balance fields are the last layer's and are hidden; good regions are not -used. +cubes, as the server's 2D training does), with Adam. Its controls are the +epochs, its own learning rates (1e-4 by default, a hundredth of the last +layer's), the server's *Batch Size*, in tiles per step (1 by default), *Tile +Size*, *Augmentation*, and the last layer's loss, distillation, smoothing and +balance fields, which it reads as the last layer does. Distillation holds the +model to the starting network's output on the voxels its scope picks in each +tile, worked out on the same tile it trains on. Good regions are not used. + +*Augmentation*, off by default as on the server, gives each step's tiles a +random brightness (×0.8 to ×1.2) and noise (1% of the raw's range) before +normalization, then random flips and quarter turns of the image and the paint +together, the server's augmentation. The targets, and the starting network's +output for distillation, are made after the turn, from the turned paint and +image, so an affinity still means its own direction. The whole network sees +the same tiles every epoch, so over many epochs it can learn them by heart: +augmentation is what keeps it from that. The last layer trains on features +computed once, and takes no augmentation. *Tile Size* is the tile's side in pixels: 172, 300 (the default), 428 or 556. The model reads 88 pixels of context around what it predicts, so a tile's diff --git a/web/ARCHITECTURE.md b/web/ARCHITECTURE.md index 7d9dd995..573c0b41 100644 --- a/web/ARCHITECTURE.md +++ b/web/ARCHITECTURE.md @@ -408,9 +408,10 @@ web/ for picked directories; paint-store.ts + paint-route.ts: the annotation volumes' files and the worker's paint/ route finetune/ types.ts (the seams), volume.ts (planVolume, the zarr v2 files), samples.ts, targets.ts, onnx-proto.ts, onnx-head.ts, features.ts, head-trainer.ts, distillation.ts, good-regions.ts, run.ts (runFinetune), - zip.ts (the paint's download), modes.ts (what a manifest offers) + zip.ts (the paint's download), modes.ts (what a manifest offers), augment.ts (the whole network's) full/ the whole network of a 2D model in TF.js: cellmap-tfjs-train.ts (Marwan's runtime), graph.ts (the - training graph, its tile sides), files.ts, tiles.ts, tfjs.ts (lazy TF.js), run-full.ts + training graph, its tile sides), files.ts, tiles.ts, tfjs.ts (lazy TF.js), loss.ts (the head's + loss in TF.js), run-full.ts ai-annotate/ AI-assisted annotation (above): types.ts, geometry.ts, plane.ts, mask.ts, labels.ts, fill.ts, prompts.ts, organelles.ts (the pure ports), images.ts (the canvas codec), gemini.ts, fake.ts, providers.ts, credentials.ts, usage.ts, audit.ts; labels.ts and fill.ts are Label it's too diff --git a/web/src/backend/routes/finetune.ts b/web/src/backend/routes/finetune.ts index c55ff4e9..7ae8cf77 100644 --- a/web/src/backend/routes/finetune.ts +++ b/web/src/backend/routes/finetune.ts @@ -12,15 +12,15 @@ // for the models directory; a run in the tab (src/finetune/run.ts) for the // cluster job, so a submit needs no output path, queue or charge group, and // the server-only fields the form still sends (LoRA rank, patches per -// epoch, rehearsal, augmentation) are ignored. What trains is what the -// model's manifest offers (finetune/modes.ts, answered as finetune_modes) -// and the form's Training field picks: the head, with the server's -// per-voxel losses (BCE, MSE, margin), its distillation term (λ and scope) -// and the browser's own batch of voxels per step (readTraining), Dice -// refused since it needs a whole patch per sample (refuseLoss); or the -// whole network of a 2D model (finetune/full), with masked BCE, the -// server's Batch Size in tiles and a tile side of its own, which the -// models route offers with what a step costs (fullTraining). The good +// epoch, rehearsal) are ignored. What trains is what the model's manifest +// offers (finetune/modes.ts, answered as finetune_modes) and the form's +// Training field picks, with the server's per-voxel losses (BCE, MSE, +// margin), smoothing, class balance and distillation term (λ and scope) in +// both (readTraining), Dice refused since it needs a whole patch per sample +// (refuseLoss): the head, with the browser's own batch of voxels per step; +// or the whole network of a 2D model (finetune/full), with the server's +// Batch Size in tiles, a tile side of its own, which the models route +// offers with what a step costs (fullTraining), and augmentation. The good // regions' routes are routes/good-regions.ts, and the patch tools' Label it // and AI annotation have routes of their own // (view-labels.ts, ai-annotate.ts); crops import keeps needing a server and @@ -543,14 +543,14 @@ function refuseLoss(loss: unknown): string | null { * trainingParams, as job-monitor.js sends them): the ones the browser's * trainer takes, checked, as run.ts's options and as the job's params (the * card shows the loss, margin and distillation from them). finetune_mode - * says what trains, "head" when absent. For the head, batch_pixels is the - * browser's own: voxels per Adam step, where the server's batch_size - * counts patches. The margin goes with the margin loss alone, and label - * smoothing is 0 with it, as the server sets it. The fields only a cluster - * job uses (lora_r, batch_size, patches_per_epoch, rehearsal_fraction, - * augment, queue, checkpoint_path) are left out of both. For the whole - * network (readFullTraining) the batch is batch_size, in tiles, and the - * tile's side tile_size; the head's fields are left out. */ + * says what trains, "head" when absent. Both modes take the loss fields + * (readLoss). For the head, batch_pixels is the browser's own: voxels per + * Adam step, where the server's batch_size counts patches. The fields only + * a cluster job uses (lora_r, patches_per_epoch, rehearsal_fraction, queue, + * checkpoint_path, and for the head batch_size and augment) are left out + * of both. For the whole network (readFullTraining) the batch is + * batch_size, in tiles, the tile's side tile_size, and augment says + * whether its tiles are augmented. */ export function readTraining(body: Record): { options: FinetuneOptions; params: Record } | { error: string } { const mode = body.finetune_mode === undefined || body.finetune_mode === null || body.finetune_mode === "" ? "head" : body.finetune_mode; if (mode !== "head" && mode !== "full") return { error: `finetune_mode: expected head or full, not ${JSON.stringify(mode)}` }; @@ -558,7 +558,23 @@ export function readTraining(body: Record): { options: Finetune if (!Number.isInteger(epochs) || epochs < 1) return { error: "num_epochs: expected a whole number of epochs, 1 or more" }; const learningRate = numberOr(body.learning_rate, mode === "full" ? FULL_LEARNING_RATE : DEFAULT_OPTIONS.learningRate); if (!Number.isFinite(learningRate) || learningRate <= 0) return { error: "learning_rate: expected a positive number" }; - if (mode === "full") return readFullTraining(body, epochs, learningRate); + const loss = readLoss(body); + if ("error" in loss) return loss; + if (mode === "full") return readFullTraining(body, epochs, learningRate, loss); + const batchPixels = numberOr(body.batch_pixels, DEFAULT_OPTIONS.batchPixels); + if (!Number.isInteger(batchPixels) || batchPixels < 1 || batchPixels > MAX_BATCH_PIXELS) return { error: `batch_pixels: expected a whole number of voxels per step, from 1 to ${MAX_BATCH_PIXELS}` }; + return { + options: { ...DEFAULT_OPTIONS, epochs, learningRate, ...loss.options, batchPixels }, + params: { num_epochs: epochs, learning_rate: learningRate, ...loss.params, batch_pixels: batchPixels, finetune_mode: "head" }, + }; +} + +type LossFields = Pick; + +/** The loss fields both modes train with, checked: the loss (Dice refused), + * its margin, label smoothing (0 with the margin loss), class balance and + * the distillation term's λ and scope. */ +function readLoss(body: Record): { options: LossFields; params: Record } | { error: string } { const lossType = body.loss_type === undefined || body.loss_type === null || body.loss_type === "" ? DEFAULT_OPTIONS.lossType : (body.loss_type as HeadLoss); if (!HEAD_LOSSES.includes(lossType)) return { error: refuseLoss(lossType)! }; const margin = numberOr(body.margin, DEFAULT_OPTIONS.margin); @@ -567,47 +583,43 @@ export function readTraining(body: Record): { options: Finetune if (!(labelSmoothing >= 0 && labelSmoothing < 1)) return { error: "label_smoothing: expected a number from 0 up to, not including, 1" }; if (lossType === "margin") labelSmoothing = 0; const balanceClasses = typeof body.balance_classes === "boolean" ? body.balance_classes : DEFAULT_OPTIONS.balanceClasses; - const batchPixels = numberOr(body.batch_pixels, DEFAULT_OPTIONS.batchPixels); - if (!Number.isInteger(batchPixels) || batchPixels < 1 || batchPixels > MAX_BATCH_PIXELS) return { error: `batch_pixels: expected a whole number of voxels per step, from 1 to ${MAX_BATCH_PIXELS}` }; const distillationLambda = numberOr(body.distillation_lambda, DEFAULT_OPTIONS.distillationLambda); if (!Number.isFinite(distillationLambda) || distillationLambda < 0) return { error: "distillation_lambda: expected a number, 0 or more (0 switches distillation off)" }; const scope = body.distillation_scope === undefined || body.distillation_scope === null || body.distillation_scope === "" ? DEFAULT_OPTIONS.distillationScope : String(body.distillation_scope).toLowerCase(); if (scope !== "unlabeled" && scope !== "all") return { error: "distillation_scope: expected unlabeled or all" }; const distillationScope = scope as DistillationScope; return { - options: { ...DEFAULT_OPTIONS, epochs, learningRate, lossType, margin, balanceClasses, labelSmoothing, batchPixels, distillationLambda, distillationScope }, + options: { lossType, margin, balanceClasses, labelSmoothing, distillationLambda, distillationScope }, params: { - num_epochs: epochs, - learning_rate: learningRate, loss_type: lossType, ...(lossType === "margin" ? { margin } : {}), balance_classes: balanceClasses, label_smoothing: labelSmoothing, - batch_pixels: batchPixels, distillation_lambda: distillationLambda, distillation_scope: distillationScope, - finetune_mode: "head", }, }; } -/** readTraining's whole-network part: the batch in tiles (batch_size, 1 to - * MAX_BATCH_PATCHES) and the tile side (tile_size, a whole number; blank - * for the default), which submit then checks against the model's graph. - * The loss is masked BCE, so the params say "bce". */ -function readFullTraining(body: Record, epochs: number, learningRate: number): { options: FinetuneOptions; params: Record } | { error: string } { +/** readTraining's whole-network part: the loss fields as the head's, the + * batch in tiles (batch_size, 1 to MAX_BATCH_PATCHES), the tile side + * (tile_size, a whole number; blank for the default), which submit then + * checks against the model's graph, and augmentation (augment). */ +function readFullTraining(body: Record, epochs: number, learningRate: number, loss: { options: LossFields; params: Record }): { options: FinetuneOptions; params: Record } | { error: string } { const batchPatches = numberOr(body.batch_size, DEFAULT_OPTIONS.batchPatches); if (!Number.isInteger(batchPatches) || batchPatches < 1 || batchPatches > MAX_BATCH_PATCHES) return { error: `batch_size: expected a whole number of tiles per step, from 1 to ${MAX_BATCH_PATCHES}` }; const tile = body.tile_size === undefined || body.tile_size === null || body.tile_size === "" ? null : Number(body.tile_size); if (tile !== null && (!Number.isInteger(tile) || tile < 1)) return { error: "tile_size: expected a whole number of pixels on a side" }; + const augment = typeof body.augment === "boolean" ? body.augment : DEFAULT_OPTIONS.augment; return { - options: { ...DEFAULT_OPTIONS, mode: "full", epochs, learningRate, batchPatches, tileSize: tile, distillationLambda: 0 }, + options: { ...DEFAULT_OPTIONS, mode: "full", epochs, learningRate, ...loss.options, batchPatches, tileSize: tile, augment }, params: { num_epochs: epochs, learning_rate: learningRate, - loss_type: "bce", + ...loss.params, batch_size: batchPatches, ...(tile === null ? {} : { tile_size: tile }), + augment, finetune_mode: "full", }, }; diff --git a/web/src/finetune/augment.ts b/web/src/finetune/augment.ts new file mode 100644 index 00000000..4d1dfac2 --- /dev/null +++ b/web/src/finetune/augment.ts @@ -0,0 +1,135 @@ +// Training-time augmentation of a whole-network tile, the port of the +// server's (cellmap_flow/finetune/data/augment.py): a random brightness +// (×0.8 to ×1.2) and Gaussian noise of 1% of the value range on the raw as +// stored, before it is normalized, then random flips of each axis and a +// random quarter turn in the y–x plane of the raw and the paint alike. +// +// The targets are made from the paint after it is turned, never turned +// themselves: an affinity's channels name directions (a flip makes the left +// neighbour the right one, a quarter turn makes y into x), so turning a +// target would need its channels permuted and shifted, while the paint's +// labels mean the same wherever they lie. The same goes for the starting +// network's logits that distillation holds the model to: they are worked +// out on the turned input, not turned. So the run turns the raw and the +// paint, then makes both from them (full/run-full.ts). +// +// A tile is a block of (z, y, x) in the order the network takes it: one +// slice for a 2D model (its slice axis first, as tiles.ts hands it over). +// The raw and the paint share the block's centre: the raw is the paint's +// box grown by the same context on each side, so the same flip or turn of +// each about its own centre keeps them registered. z is flipped but never +// turned into, as on the server: EM is often coarser along z. +import type { Step } from "../contracts"; + +/** One tile's draw: a flip per axis (z, y, x), the quarter turns in y–x (0 to 3), the brightness factor. */ +export interface AugmentDraw { + flips: [boolean, boolean, boolean]; + turns: number; + scale: number; +} + +/** A draw from `random`: each axis flipped with even odds, 0 to 3 quarter + * turns when the y–x plane may be turned (`turnable`: square, and its two + * voxel sizes the same), and a brightness factor from 0.8 to 1.2. */ +export function drawAugmentation(random: () => number, turnable: boolean): AugmentDraw { + const flips: [boolean, boolean, boolean] = [random() < 0.5, random() < 0.5, random() < 0.5]; + const turns = turnable ? Math.min(3, Math.floor(random() * 4)) : 0; + return { flips, turns, scale: 0.8 + 0.4 * random() }; +} + +type Values = Float32Array | Uint8Array | Uint16Array; + +/** `data`, a block of `shape` (z, y, x, C order), flipped along each axis + * `draw` flips, then turned `draw.turns` quarter turns in y–x (which needs + * y and x the same length): a new array of the same type. */ +export function transformBlock(data: T, shape: readonly [number, number, number], draw: AugmentDraw): T { + const [d, h, w] = shape; + if (data.length !== d * h * w) throw new Error(`a block of ${shape.join("×")} has ${d * h * w} values, not ${data.length}`); + const turns = ((draw.turns % 4) + 4) % 4; + if (turns && h !== w) throw new Error(`a quarter turn needs a square y–x plane, not ${h}×${w}`); + const [fz, fy, fx] = draw.flips; + const out = new (data.constructor as new (n: number) => T)(data.length); + for (let z = 0; z < d; z++) { + const sz = fz ? d - 1 - z : z; + for (let y = 0; y < h; y++) for (let x = 0; x < w; x++) { + // The voxel of the flipped block that lands at (y, x) after the turns. + let a = y, b = x; + for (let k = 0; k < turns; k++) { const t = a; a = b; b = h - 1 - t; } + const sy = fy ? h - 1 - a : a, sx = fx ? w - 1 - b : b; + out[(z * h + y) * w + x] = data[(sz * h + sy) * w + sx]; + } + } + return out; +} + +const INTEGER_RANGES: Record = { + uint8: [0, 255], int8: [-128, 127], uint16: [0, 65535], int16: [-32768, 32767], + uint32: [0, 4294967295], int32: [-2147483648, 2147483647], +}; + +/** The range the raw is kept inside (augment.py _intensity_range): the first + * normalizer's window when it is a MinMaxNormalizer, the range the data is + * known to use and is clipped to next anyway; else an integer dtype's own; + * else none (a float raw, which keeps its negative half). */ +export function intensityRange(norm: readonly Step[], dtype: string): [number, number] | null { + const first = norm[0]; + if (first?.name === "MinMaxNormalizer") { + const low = Number(first.min_value ?? 0), high = Number(first.max_value ?? 255); + if (high > low) return [low, high]; + } + return INTEGER_RANGES[dtype] ?? null; +} + +/** A standard normal draw from `random` (Box–Muller). */ +function gaussian(random: () => number): number { + const u = 1 - random(), v = random(); + return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v); +} + +/** The raw voxels `inside` the data scaled by `scale`, plus Gaussian noise of + * 1% of `range` (of the block's largest magnitude when there is none), + * clipped to `range`; 0 outside, where the model stage pads with zeros. */ +export function augmentIntensity(raw: Float32Array, inside: Uint8Array, scale: number, range: [number, number] | null, random: () => number): Float32Array { + let span: number; + if (range) { + span = range[1] - range[0]; + } else { + let largest = 0; + for (let i = 0; i < raw.length; i++) if (inside[i]) largest = Math.max(largest, Math.abs(raw[i])); + span = Math.max(1, largest); + } + const sigma = 0.01 * span, out = new Float32Array(raw.length); + for (let i = 0; i < raw.length; i++) { + if (!inside[i]) continue; + const v = raw[i] * scale + sigma * gaussian(random); + out[i] = range ? Math.min(Math.max(v, range[0]), range[1]) : v; + } + return out; +} + +/** What a run's draws were, for its log: proof that augmentation ran, and ran as drawn. */ +export class AugmentTally { + private n = 0; + private flipped = [0, 0, 0]; + private turned = [0, 0, 0, 0]; + private low = Infinity; + private high = -Infinity; + + add(draw: AugmentDraw): void { + this.n++; + draw.flips.forEach((f, a) => { if (f) this.flipped[a]++; }); + this.turned[draw.turns]++; + this.low = Math.min(this.low, draw.scale); + this.high = Math.max(this.high, draw.scale); + } + + /** "120 tiles augmented: flipped along rows 49%, columns 52%; turned 0°: 26%, …; brightness ×0.80 to ×1.20". + * `names` name the block's axes (z, y, x), null for one a tile is a voxel thick along (a slice's z). */ + describe(names: readonly (string | null)[], turnable: boolean): string { + if (!this.n) return "no tiles augmented"; + const pct = (k: number) => `${Math.round((100 * k) / this.n)}%`; + const flips = names.map((name, a) => (name ? `${name} ${pct(this.flipped[a])}` : "")).filter(Boolean).join(", "); + const turns = turnable ? `; turned ${this.turned.map((k, i) => `${i * 90}°: ${pct(k)}`).join(", ")}` : ""; + return `${this.n} tile${this.n === 1 ? "" : "s"} augmented: flipped along ${flips}${turns}; brightness ×${this.low.toFixed(2)} to ×${this.high.toFixed(2)}`; + } +} diff --git a/web/src/finetune/full/loss.ts b/web/src/finetune/full/loss.ts new file mode 100644 index 00000000..863b85ae --- /dev/null +++ b/web/src/finetune/full/loss.ts @@ -0,0 +1,75 @@ +// The whole network's loss in TF.js: the same terms the head trains with +// (head-trainer.ts batchLoss and distillationTerm, the port of the server's +// losses.py), so a whole-network run and a head run on the same paint are +// judged alike, as the server's runs are. +// +// - The supervised term on the painted voxels: BCE on the logits against +// the smoothed target, the squared distance of the probability from it +// (MSE), or the margin loss; reduced over the mask as one mean, or with +// class balance as the mean of the foreground's and the background's +// means, split by the target before smoothing. +// - λ × the distillation term: the mean squared distance of the logits from +// the starting network's on the same (augmented) input, over the voxels +// the scope picks: those the mask leaves out (unlabeled), or all. +// +// As on the server (lora_trainer.py), the starting network sees the tile +// the student sees, after augmentation; its logits are not turned. +import type * as TF from "@tensorflow/tfjs-core"; + +import type { FinetuneOptions } from "../types"; +import type { Tf, TrainNet } from "./cellmap-tfjs-train"; + +export type FullLossOptions = Pick; + +/** Σ l·w / max(Σ w, 1): masked_mean, whose clamp makes a set that weighs nothing cost nothing. */ +function weightedMean(tf: Tf, l: TF.Tensor, w: TF.Tensor): TF.Scalar { + return tf.div(tf.sum(tf.mul(l, w)), tf.maximum(tf.sum(w), 1)) as TF.Scalar; +} + +/** The supervised term of logits `z` against the 0/1 targets `y` over the mask `m` (batchLoss). */ +export function supervisedLoss(tf: Tf, z: TF.Tensor, y: TF.Tensor, m: TF.Tensor, o: FullLossOptions): TF.Scalar { + const s = o.lossType === "margin" ? 0 : o.labelSmoothing; + const target = s ? tf.add(tf.mul(y, 1 - s), s / 2) : y; + let l: TF.Tensor; + if (o.lossType === "mse") { + l = tf.square(tf.sub(tf.sigmoid(z), target)); + } else if (o.lossType === "margin") { + const p = tf.sigmoid(z); + const short = tf.add(tf.mul(y, tf.relu(tf.sub(1 - o.margin, p))), tf.mul(tf.sub(1, y), tf.relu(tf.sub(p, o.margin)))); + l = tf.square(short); + } else { + // max(z, 0) − z·y' + log(1 + exp(−|z|)): BCE on logits that overflows for no z. + l = tf.add(tf.sub(tf.relu(z), tf.mul(z, target)), tf.log1p(tf.exp(tf.neg(tf.abs(z))))); + } + if (!o.balanceClasses) return weightedMean(tf, l, m); + const fg = tf.mul(m, y), bg = tf.mul(m, tf.sub(1, y)); + return tf.mul(tf.add(weightedMean(tf, l, fg), weightedMean(tf, l, bg)), 0.5) as TF.Scalar; +} + +/** The distillation term (distillationTerm): the weighted mean of (z − z₀)², weights `w`. */ +export function distillationLoss(tf: Tf, z: TF.Tensor, teacher: TF.Tensor, w: TF.Tensor): TF.Scalar { + return weightedMean(tf, tf.square(tf.sub(z, teacher)), w); +} + +/** The distillation weights of mask `m` for `scope`: 1 − m (unlabeled) or 1 everywhere (all). */ +export function distillationWeights(tf: Tf, m: TF.Tensor, scope: FullLossOptions["distillationScope"]): TF.Tensor { + return scope === "all" ? tf.onesLike(m) : tf.sub(1, m); +} + +/** One Adam step of `net` on the batch: the supervised term, plus λ × the + * distillation term towards `teacher`'s logits when there are some. Both + * terms as they were before the step; the caller disposes them. */ +export function fullTrainStep(tf: Tf, net: TrainNet, opt: TF.Optimizer, x: TF.Tensor4D, y: TF.Tensor4D, m: TF.Tensor4D, teacher: TF.Tensor4D | null, o: FullLossOptions): { total: TF.Scalar; supervised: TF.Scalar; distillation: TF.Scalar | null } { + let supervised: TF.Scalar | null = null, distillation: TF.Scalar | null = null; + const total = opt.minimize(() => { + const z = net.forward(x); + const sup = tf.keep(supervisedLoss(tf, z, y, m, o)); + supervised = sup; + if (!teacher || !(o.distillationLambda > 0)) return tf.clone(sup); + const dist = tf.keep(distillationLoss(tf, z, teacher, distillationWeights(tf, m, o.distillationScope))); + distillation = dist; + return tf.add(sup, tf.mul(dist, o.distillationLambda)) as TF.Scalar; + }, true, net.vars); + if (!total || !supervised) throw new Error("the optimizer returned no loss"); + return { total, supervised, distillation }; +} diff --git a/web/src/finetune/full/run-full.ts b/web/src/finetune/full/run-full.ts index a536712c..e5292676 100644 --- a/web/src/finetune/full/run-full.ts +++ b/web/src/finetune/full/run-full.ts @@ -2,34 +2,47 @@ // training graph and weight offsets fetched (files.ts), the paint cut into // tiles in the planes it was drawn in and read with its raw (tiles.ts), // TF.js loaded (tfjs.ts), the network built from the weights in the served -// .onnx and trained on the tiles with masked BCE on its logits and Adam -// (Marwan Zouinkhi's runtime, cellmap-tfjs-train.ts), and the trained -// weights written back into a copy of the .onnx at their offsets, so the -// result serves in onnxruntime-web exactly as the model did. +// .onnx (Marwan Zouinkhi's runtime, cellmap-tfjs-train.ts) and trained on +// the tiles with Adam, and the trained weights written back into a copy of +// the .onnx at their offsets, so the result serves in onnxruntime-web +// exactly as the model did. // // What trains: every param the graph marks trainable but the manifest's // `frozen` (BatchNorm's running statistics, which the forward pass reads as -// constants), on the painted voxels alone: no distillation, no good -// regions, no class balance or label smoothing, which are the head's. -// The epochs and the learning rate are the form's; the batch is in tiles -// (the form's Batch Size), the tile's side the form's Tile size, checked -// against the sides the graph takes (graph.ts). +// constants). The loss is the head's (loss.ts): the form's supervised loss +// on the painted voxels, smoothed and class-balanced as it says, plus λ × +// the distance of the logits from the starting network's (the base's +// weights, also on a restart) over the voxels its scope picks. Good regions +// are not used. The epochs and the learning rate are the form's; the batch +// is in tiles (the form's Batch Size), the tile's side the form's Tile +// size, checked against the sides the graph takes (graph.ts). +// +// With augmentation (augment.ts), each step's tiles are drawn anew from +// their raw as stored and their paint: brightness and noise on the raw, +// then the same flips and quarter turn of the raw and the paint, then the +// targets made from the turned paint and the starting network run on the +// turned tile, in the server's order (lora_trainer.py), so every target +// means what it did whatever the turn. // // The pieces are a parameter (createRunFull) so a test can run it with its // own TF.js loader and files. +import type * as TF from "@tensorflow/tfjs-core"; + import type { DatasetInfo, ModelEntry } from "../../contracts"; import { chooseRawLevel, geometryOf, inputShape, outputLevel, outputShape, ratio, sameVoxelSize } from "../../infer/geometry"; import { mulberry32, shuffle, SLICE_MS, yieldEvery } from "../head-trainer"; import { count, epochSummary, plural } from "../log-words"; +import { AugmentTally, augmentIntensity, drawAugmentation, intensityRange, transformBlock } from "../augment"; import { refuseMode, sourceOf } from "../modes"; import { readGraph, readModel } from "../onnx-proto"; import { targetTransformFor } from "../targets"; import type { Finetuned, RunFinetune } from "../types"; -import { buildNet, type InitializerOffsets, readOnnxWeights, trainStep, writeOnnxWeights } from "./cellmap-tfjs-train"; +import { buildNet, type InitializerOffsets, readOnnxWeights, writeOnnxWeights } from "./cellmap-tfjs-train"; import { type FetchFullFiles, fetchFullFiles } from "./files"; import { defaultTile, describeSides, outputSide, tileChoices, tileSides } from "./graph"; +import { fullTrainStep } from "./loss"; import { type LoadedTf, loadTf } from "./tfjs"; -import { paintDepthAxes, paintedVoxels, planTiles, PLANE_NAMES, readPaint, type TileSample, tileInput, tileLabels, tileTargets } from "./tiles"; +import { normalizeTile, paintDepthAxes, paintedVoxels, planeAxes, planTiles, PLANE_NAMES, readPaint, type TileSample, tileLabels, tileRaw, tileTargets } from "./tiles"; export interface FullDeps { loadTf: () => Promise; @@ -117,6 +130,11 @@ export function createRunFull(deps: FullDeps = { loadTf, fetchFiles: fetchFullFi const names = Object.keys(graph.params); const start = new Uint8Array(run.initialBytes ?? run.baseBytes); checkOffsets(start, offsets, names, source.file); + // The teacher is the base's weights, the model the job started from, also + // when a restart starts from the last run's (as the head's is). + const lambda = options.distillationLambda; + const base = lambda > 0 ? new Uint8Array(run.baseBytes) : null; + if (base) checkOffsets(base, offsets, names, source.file); const side = checkTileSide(graph, options.tileSize, entry.name); const output = side - graph.context; const frozen = new Set(spec.frozen); @@ -139,18 +157,34 @@ export function createRunFull(deps: FullDeps = { loadTf, fetchFiles: fetchFullFi const tiles = planTiles(voxels, depth, output); const inPlanes = planes.map((n, d) => (n ? `${count(n)} in ${PLANE_NAMES[d]} slices` : "")).filter(Boolean).join(", "); report.log(`Finetuning ${entry.name}'s whole network on volume ${volume.id}: ${count(painted)} painted voxels (${inPlanes}) in ${plural(tiles.length, "tile")} of ${side}² (output ${output}²)`); - report.log(`The network: ${count(trainableValues)} weights in ${plural(trainable.size, "tensor")} train${frozen.size ? `, ${plural(frozen.size, "frozen statistic")} held` : ""}; ${plural(options.epochs, "epoch")} of ${plural(Math.min(options.batchPatches, tiles.length), "tile")} per step at learning rate ${options.learningRate}, masked BCE on the logits, Adam`); - if (run.goodRegions?.length) report.log(`${plural(run.goodRegions.length, "good region")} marked: whole-network training does not use them (it trains on the paint alone; the last layer's training holds the model in them)`); + report.log(`The network: ${count(trainableValues)} weights in ${plural(trainable.size, "tensor")} train${frozen.size ? `, ${plural(frozen.size, "frozen statistic")} held` : ""}; ${plural(options.epochs, "epoch")} of ${plural(Math.min(options.batchPatches, tiles.length), "tile")} per step at learning rate ${options.learningRate}, ${lossWords(options)}, Adam`); + report.log(lambda > 0 + ? `Distillation, λ ${lambda}: held to the starting network on ${options.distillationScope === "all" ? "every voxel" : "the unpainted voxels"} of each tile's output` + : "Distillation off (λ 0): trained on the paint alone"); + if (run.goodRegions?.length) report.log(`${plural(run.goodRegions.length, "good region")} marked: whole-network training does not use them (the last layer's training holds the model in them)`); + const norm = volume.input_norm, range = intensityRange(norm, dataset.dtype); + // A tile turns where its plane's two voxel sizes are the same, in the raw and the model. + const turnable = (depth: number) => { + const [u, v] = planeAxes(depth as 0 | 1 | 2); + return [geo.g.input_voxel_size, geo.raw.voxel].every((size) => sameVoxelSize([size[u]], [size[v]])); + }; + report.log(options.augment + ? `Augmentation on: each step's tiles get a brightness ×0.8 to ×1.2 and noise of 1% of ${range ? `the raw's range (${range[0]} to ${range[1]})` : "each tile's largest value"} on the raw before it is normalized, then random flips and quarter turns of the raw and the paint alike, with the targets${lambda > 0 ? " and the starting network's logits" : ""} made after` + : "Augmentation off"); const transform = deps.targetTransformFor(entry); const samples: TileSample[] = []; + let distilled = 0; for (const tile of tiles) { if (report.cancelled()) throw new Error("cancelled before training"); const labels = tileLabels(tile, output, chunks, volume); const t = tileTargets(labels, output, transform); if (!t.masked) continue; // paint the loss cannot count here (an affinity needs a painted pair) - const input = await tileInput(tile, side, output, run.infer, dataset, geo.rawLevel, volume.input_norm); - samples.push({ tile, input, target: t.target, mask: t.mask, masked: t.masked }); + const { raw, inside } = await tileRaw(tile, side, output, run.infer, dataset, geo.rawLevel); + samples.push(options.augment + ? { tile, input: new Float32Array(0), target: t.target, mask: t.mask, masked: t.masked, raw, inside, labels } + : { tile, input: normalizeTile(raw, inside, norm), target: t.target, mask: t.mask, masked: t.masked }); + if (lambda > 0) distilled += options.distillationScope === "all" ? output * output : unmaskedVoxels(t.mask, t.channels); } if (!samples.length) throw new Error("the paint gives the loss nothing to count: an affinity is between two painted neighbours, so paint strokes, not single voxels"); if (samples.length < tiles.length) report.log(`${plural(tiles.length - samples.length, "tile")} of lone painted voxels give the loss nothing to count (an affinity needs two painted neighbours) and ${tiles.length - samples.length === 1 ? "is" : "are"} left out`); @@ -162,15 +196,19 @@ export function createRunFull(deps: FullDeps = { loadTf, fetchFiles: fetchFullFi report.log(backend === "webgpu" ? "TF.js on WebGPU" : `TF.js on its CPU backend: ${fallback ?? "its WebGPU backend did not start"}. It trains, but a step takes minutes here where the GPU takes a fraction of a second; reloading the page tries WebGPU again, and a smaller tile is quicker`); - let net: ReturnType | null = null, opt: ReturnType | null = null; + // Augmentation's draws: a stream of its own, so a run without it shuffles as before. + const augmentRandom = mulberry32((options.seed ^ 0x9e3779b9) >>> 0); + const tally = new AugmentTally(); + let net: ReturnType | null = null, teacher: ReturnType | null = null, opt: ReturnType | null = null; try { net = buildNet(tf, graph, readOnnxWeights(start, offsets, names), { trainable }); + if (base) teacher = buildNet(tf, graph, readOnnxWeights(base, offsets, names), { trainable: new Set() }); opt = tf.train.adam(options.learningRate); const batch = Math.max(1, Math.min(options.batchPatches, samples.length)); const order = new Uint32Array(samples.length).map((_, i) => i); const random = mulberry32(options.seed); const breathe = yieldEvery(SLICE_MS); - const losses: number[] = []; + const losses: number[] = [], supervisedLosses: number[] = []; const x = new Float32Array(batch * side * side), t = new Float32Array(batch * output * output * channels), m = new Float32Array(t.length); let steps = 0, firstStepMs = 0; const stepsPerEpoch = Math.ceil(samples.length / batch); @@ -178,12 +216,12 @@ export function createRunFull(deps: FullDeps = { loadTf, fetchFiles: fetchFullFi if (report.cancelled()) break; report.log(`Starting epoch ${epoch} of ${options.epochs}`); shuffle(order, random); - let sum = 0, weight = 0; + let sum = 0, sumSupervised = 0, sumDistillation = 0, weight = 0; for (let from = 0; from < samples.length; from += batch) { const rows = Math.min(batch, samples.length - from); let masked = 0; for (let r = 0; r < rows; r++) { - const s = samples[order[from + r]]; + const s = augmented(samples[order[from + r]]); x.set(s.input, r * s.input.length); t.set(s.target, r * s.target.length); m.set(s.mask, r * s.mask.length); @@ -194,16 +232,22 @@ export function createRunFull(deps: FullDeps = { loadTf, fetchFiles: fetchFullFi const ts = tf.tensor4d(t.slice(0, rows * output * output * channels), [rows, output, output, channels]); const ms = tf.tensor4d(m.slice(0, rows * output * output * channels), [rows, output, output, channels]); const t0 = performance.now(); - let value: number; + let value: number, supervised: number, distillation = 0; + let zs: TF.Tensor4D | null = null; try { - const loss = trainStep(tf, net, opt, xs, ts, ms); - value = (await loss.data())[0]; - loss.dispose(); + if (teacher) zs = tf.tidy(() => teacher!.forward(xs)); + const step = fullTrainStep(tf, net, opt, xs, ts, ms, zs, options); + value = (await step.total.data())[0]; + supervised = (await step.supervised.data())[0]; + if (step.distillation) distillation = (await step.distillation.data())[0]; + step.total.dispose(); step.supervised.dispose(); step.distillation?.dispose(); } finally { - xs.dispose(); ts.dispose(); ms.dispose(); + xs.dispose(); ts.dispose(); ms.dispose(); zs?.dispose(); } if (!Number.isFinite(value)) throw new Error(`the loss is ${value} at epoch ${epoch}: the learning rate ${options.learningRate} is too high for the whole network; lower it`); sum += value * masked; + sumSupervised += supervised * masked; + sumDistillation += distillation * masked; weight += masked; steps++; // The second step's time: the first compiles the GPU's shaders. @@ -216,23 +260,59 @@ export function createRunFull(deps: FullDeps = { loadTf, fetchFiles: fetchFullFi if (report.cancelled()) break training; } losses.push(sum / weight); - report.log(`Epoch ${epoch}/${options.epochs} - Loss: ${(sum / weight).toFixed(6)}`); + supervisedLosses.push(sumSupervised / weight); + report.log(teacher + ? `Epoch ${epoch}/${options.epochs} - Loss: ${(sum / weight).toFixed(6)} - Supervised: ${(sumSupervised / weight).toFixed(6)} - Distillation: ${(sumDistillation / weight).toFixed(6)}` + : `Epoch ${epoch}/${options.epochs} - Loss: ${(sum / weight).toFixed(6)}`); } + if (options.augment) report.log(tally.describe([null, "rows", "columns"], samples.some((s) => turnable(s.tile.depth)))); const bytes = writeOnnxWeights(start, offsets, net.exportWeights()); const finalLoss = losses.length ? losses[losses.length - 1] : null; - report.log(epochSummary(losses, options.epochs)); + const finalSupervised = supervisedLosses.length ? supervisedLosses[supervisedLosses.length - 1] : null; + report.log(epochSummary(losses, options.epochs, teacher && finalSupervised !== null ? ` (supervised ${finalSupervised.toFixed(6)})` : "")); return { - bytes: bytes.buffer, epochs: losses.length, losses, finalLoss, supervisedLosses: losses.slice(), - pixels: painted, chunks: chunks.size, distilled: 0, tiles: samples.length, + bytes: bytes.buffer, epochs: losses.length, losses, finalLoss, supervisedLosses, + pixels: painted, chunks: chunks.size, distilled, tiles: samples.length, }; } finally { net?.dispose(); + teacher?.dispose(); opt?.dispose(); release(); } + + /** `s` as this step trains on it: as read, or with augmentation a new + * draw of its raw and paint, normalized, turned, and its targets made + * from the turned paint. */ + function augmented(s: TileSample): Pick { + if (!s.raw || !s.inside || !s.labels) return s; + const draw = drawAugmentation(augmentRandom, turnable(s.tile.depth)); + tally.add(draw); + const input = transformBlock(normalizeTile(augmentIntensity(s.raw, s.inside, draw.scale, range, augmentRandom), s.inside, norm), [1, side, side], draw); + const t = tileTargets(transformBlock(s.labels, [1, output, output], draw), output, transform); + return { input, target: t.target, mask: t.mask, masked: t.masked }; + } }; } +/** The voxels of a tile's output that some channel's mask leaves out: what "unlabeled" distillation holds. */ +function unmaskedVoxels(mask: Uint8Array, channels: number): number { + let n = 0; + for (let p = 0; p < mask.length; p += channels) { + let all = true; + for (let c = 0; c < channels; c++) if (!mask[p + c]) { all = false; break; } + if (!all) n++; + } + return n; +} + +/** The supervised loss as the log names it: "BCE on the logits, class-balanced, smoothed 0.1". */ +function lossWords(o: { lossType: string; margin: number; labelSmoothing: number; balanceClasses: boolean }): string { + const name = o.lossType === "mse" ? "MSE on the probabilities" : o.lossType === "margin" ? `the margin loss (margin ${o.margin})` : "BCE on the logits"; + const smoothing = o.lossType !== "margin" && o.labelSmoothing > 0 ? `, smoothed ${o.labelSmoothing}` : ""; + return `${name} on the painted voxels${o.balanceClasses ? ", class-balanced" : ""}${smoothing}`; +} + /** "12 s", "3.4 min". */ function seconds(s: number): string { return s < 90 ? `${s < 10 ? s.toFixed(1) : Math.round(s)} s` : `${(s / 60).toFixed(1)} min`; diff --git a/web/src/finetune/full/tiles.ts b/web/src/finetune/full/tiles.ts index bc830475..039c00dd 100644 --- a/web/src/finetune/full/tiles.ts +++ b/web/src/finetune/full/tiles.ts @@ -52,6 +52,12 @@ export interface TileSample { mask: Uint8Array; /** Mask entries set, over all channels: what the loss averages over. */ masked: number; + /** Kept for augmentation, which draws each step's input and targets from + * them (run-full.ts): the raw as stored and which of it is inside the + * level (tileRaw), and the paint (tileLabels). */ + raw?: Float32Array; + inside?: Uint8Array; + labels?: Uint8Array | Uint16Array; } /** The two axes of the plane cut across `depth`, in the order a tile hands them over. */ @@ -191,12 +197,11 @@ export function tileLabels(tile: Tile, output: number, chunks: Map, dataset: DatasetInfo, level: number, norm: Step[]): Promise { +/** The raw slice around `tile`'s output box (`side` × `side`, the box grown + * by half the context each way) on raw level `level`, as stored, and which + * of its voxels are inside the level (1) or past it (0, where the value is + * 0). The raw level's voxels are the volume's (a ratio of 1). */ +export async function tileRaw(tile: Tile, side: number, output: number, infer: Pick, dataset: DatasetInfo, level: number): Promise<{ raw: Float32Array; inside: Uint8Array }> { const halo = (side - output) / 2; const [pu, pv] = planeAxes(tile.depth); const lo = [0, 0, 0], hi = [0, 0, 0]; @@ -204,17 +209,40 @@ export async function tileInput(tile: Tile, side: number, output: number, infer: lo[pu] = tile.corner[0] - halo; hi[pu] = tile.corner[0] + output + halo; lo[pv] = tile.corner[1] - halo; hi[pv] = tile.corner[1] + output + halo; const shape = hi.map((h, a) => h - lo[a]); - const block = new Float32Array(side * side); + const raw = new Float32Array(side * side), inside = new Uint8Array(side * side); const rawShape = dataset.levels[level].shape; const readLo = lo.map((o) => Math.max(o, 0)), readHi = hi.map((h, a) => Math.min(h, rawShape[a])); const span = readHi.map((h, a) => h - readLo[a]); - if (span.some((n) => n <= 0)) return block; + if (span.some((n) => n <= 0)) return { raw, inside }; const read = padded(await infer.readRaw(level, readLo, readHi), dataset.dtype, 1, readLo, readHi, readLo, readHi, "zero"); - const values = applyNormalizers(Float32Array.from(read.data as ArrayLike), norm); // The box with its depth axis one voxel thick, in C order, is the tile's // rows and columns: copied into the block where it lies in the box. - copyBox(block, shape, readLo.map((o, a) => o - lo[a]), values, span, [0, 0, 0], span); - return block; + const at = readLo.map((o, a) => o - lo[a]); + copyBox(raw, shape, at, Float32Array.from(read.data as ArrayLike), span, [0, 0, 0], span); + for (let z = 0; z < span[0]; z++) for (let y = 0; y < span[1]; y++) { + const from = ((z + at[0]) * shape[1] + y + at[1]) * shape[2] + at[2]; + inside.fill(1, from, from + span[2]); + } + return { raw, inside }; +} + +/** `raw`'s voxels `inside` normalized through `norm`, in order, as the model + * stage normalizes a block, and those past the level zero, as it pads one + * (infer/stages.ts modelInput). */ +export function normalizeTile(raw: Float32Array, inside: Uint8Array, norm: Step[]): Float32Array { + let n = 0; + for (let i = 0; i < inside.length; i++) n += inside[i]; + const values = new Float32Array(n); + for (let i = 0, k = 0; i < raw.length; i++) if (inside[i]) values[k++] = raw[i]; + const normalized = applyNormalizers(values, norm), out = new Float32Array(raw.length); + for (let i = 0, k = 0; i < raw.length; i++) if (inside[i]) out[i] = normalized[k++]; + return out; +} + +/** The network's input for `tile`: tileRaw normalized (normalizeTile). */ +export async function tileInput(tile: Tile, side: number, output: number, infer: Pick, dataset: DatasetInfo, level: number, norm: Step[]): Promise { + const { raw, inside } = await tileRaw(tile, side, output, infer, dataset, level); + return normalizeTile(raw, inside, norm); } /** A tile's targets from its labels, through the model's transform diff --git a/web/src/finetune/types.ts b/web/src/finetune/types.ts index 56e686e3..8fdb63c1 100644 --- a/web/src/finetune/types.ts +++ b/web/src/finetune/types.ts @@ -20,6 +20,7 @@ // with the volume, and where they fall on its grid; // run.ts all of it, one run: runFinetune; // modes.ts the ways a model can be trained, from its manifest; +// augment.ts the whole network's augmentation of a tile; // full/ the other way, the whole network of a 2D model trained // with TensorFlow.js (full/run-full.ts), when the // manifest offers it. @@ -122,8 +123,8 @@ export type DistillationScope = "unlabeled" | "all"; * /api/finetune/submit and the routes read them. */ export interface FinetuneOptions { /** What trains: the head, or the whole network (full/run-full.ts), which - * reads the epochs, the learning rate, the seed and its two fields below - * and none of the head's (its loss is masked BCE on the paint alone). */ + * reads the same loss, smoothing, balance and distillation fields, its + * own batch and tile below, and augmentation, but not batchPixels. */ mode: FinetuneMode; epochs: number; learningRate: number; @@ -151,6 +152,10 @@ export interface FinetuneOptions { /** The whole network's tile side, in pixels of the input; null for the * default (full/graph.ts DEFAULT_TILE, the side the graph takes nearest it). */ tileSize: number | null; + /** Whole network only: each step's tiles get a random brightness and + * noise, flips and quarter turns, as the server's augmentation + * (augment.ts); the head trains on features computed once, and ignores it. */ + augment: boolean; } /** How a run speaks to the job that follows it: its log lines, in the words @@ -183,7 +188,9 @@ export interface Finetuned { /** The painted voxels trained on, over all channels' masks, and the chunks they came from. */ pixels: number; chunks: number; - /** The voxels the distillation term held to the starting head; 0 without distillation. */ + /** The voxels the distillation term held to the starting head (the + * starting network, after a whole-network run: each epoch's); 0 without + * distillation. */ distilled: number; /** The tiles a whole-network run trained on (full/tiles.ts). */ tiles?: number; @@ -272,4 +279,5 @@ export const DEFAULT_OPTIONS: FinetuneOptions = { seed: 0, batchPatches: 1, tileSize: null, + augment: false, }; diff --git a/web/test/unit/backend-finetune.test.ts b/web/test/unit/backend-finetune.test.ts index 8d16a8df..5a0a1b7b 100644 --- a/web/test/unit/backend-finetune.test.ts +++ b/web/test/unit/backend-finetune.test.ts @@ -434,7 +434,7 @@ describe("a job's life", () => { expect(run.infer).toBe(ctx.infer); expect(run.options).toEqual({ mode: "head", epochs: 3, learningRate: 0.01, lossType: "bce", margin: 0.3, balanceClasses: true, labelSmoothing: 0, batchPixels: 4096, - distillationLambda: 0.01, distillationScope: "unlabeled", seed: 0, batchPatches: 1, tileSize: null, + distillationLambda: 0.01, distillationScope: "unlabeled", seed: 0, batchPatches: 1, tileSize: null, augment: false, }); expect(run.goodRegions).toBeUndefined(); expect(run.initialHead).toBeUndefined(); diff --git a/web/test/unit/finetune-augment.test.ts b/web/test/unit/finetune-augment.test.ts new file mode 100644 index 00000000..19fa65e7 --- /dev/null +++ b/web/test/unit/finetune-augment.test.ts @@ -0,0 +1,94 @@ +// Whole-network augmentation and loss (src/finetune/augment.ts, full/loss.ts): +// the raw and the paint stay registered through a flip and a turn, targets +// made after the turn mean what they did, intensity stays in the raw's +// range, and the TF.js loss is the head's (head-trainer.ts) term for term. +import { describe, expect, it } from "vitest"; + +import { readTraining } from "../../src/backend/routes/finetune"; +import { augmentIntensity, intensityRange, transformBlock } from "../../src/finetune/augment"; +import { distillationLoss, distillationWeights, supervisedLoss } from "../../src/finetune/full/loss"; +import { loadTf } from "../../src/finetune/full/tfjs"; +import { tileTargets } from "../../src/finetune/full/tiles"; +import { batchLoss, mulberry32 } from "../../src/finetune/head-trainer"; +import { affinityOffsets, affinityTargets } from "../../src/finetune/targets"; + +describe("augmentation of a tile", () => { + it("turns the raw and the paint alike, about their shared centre", () => { + // A 6² raw around a 2² output (a halo of 2): a mark on the output's top left voxel. + const side = 6, output = 2, halo = 2; + const raw = new Float32Array(side * side), labels = new Uint8Array(output * output); + raw[halo * side + halo] = 1; + labels[0] = 1; + for (const draw of [{ flips: [false, true, false], turns: 0 }, { flips: [false, false, true], turns: 1 }, { flips: [true, true, true], turns: 3 }] as const) { + const d = { flips: [...draw.flips] as [boolean, boolean, boolean], turns: draw.turns, scale: 1 }; + const r = transformBlock(raw, [1, side, side], d), l = transformBlock(labels, [1, output, output], d); + const at = r.indexOf(1), lt = l.indexOf(1); + expect([Math.floor(at / side) - halo, (at % side) - halo]).toEqual([Math.floor(lt / output), lt % output]); + } + }); + + it("makes affinities from the turned paint, so a boundary across x becomes one across y", () => { + // Two objects side by side (labels 2 | 3): a boundary between columns 1 and 2. + const n = 4, labels = new Uint8Array(n * n).map((_, i) => (i % n < 2 ? 2 : 3)); + const transform = affinityTargets(affinityOffsets(["mito_aff_z", "mito_aff_y", "mito_aff_x"]), 3); + // The painted pairs across a boundary on channel c: masked, target 0. + const cut = (t: ReturnType, c: number) => Array.from({ length: n * n }, (_, p) => t.mask[p * 3 + c] && !t.target[p * 3 + c]).filter(Boolean).length; + const before = tileTargets(labels, n, transform); + const after = tileTargets(transformBlock(labels, [1, n, n], { flips: [false, false, false], turns: 1, scale: 1 }), n, transform); + expect(cut(before, 2)).toBeGreaterThan(0); + expect(cut(before, 1)).toBe(0); + expect([cut(after, 1), cut(after, 2)]).toEqual([cut(before, 2), 0]); + }); + + it("keeps brightness and noise inside the raw's range, and the padding at zero", () => { + const raw = new Float32Array(1000).map((_, i) => (i * 7) % 256), inside = new Uint8Array(1000).fill(1); + inside.fill(0, 900); + const range = intensityRange([{ name: "MinMaxNormalizer", min_value: 0, max_value: 255 }], "uint8"); + expect(range).toEqual([0, 255]); + expect(intensityRange([], "uint16")).toEqual([0, 65535]); + expect(intensityRange([], "float32")).toBeNull(); + const out = augmentIntensity(raw, inside, 1.2, range, mulberry32(1)); + expect(Math.min(...out)).toBeGreaterThanOrEqual(0); + expect(Math.max(...out)).toBeLessThanOrEqual(255); + expect(out.subarray(900).every((v) => v === 0)).toBe(true); + // Brighter on average, by about the factor. + const mean = (a: ArrayLike) => Array.from(a).slice(0, 900).reduce((s, v) => s + v, 0) / 900; + expect(mean(out) / mean(raw)).toBeGreaterThan(1.1); + }); +}); + +describe("the whole network's loss", () => { + it("is the head's supervised term for every loss, balanced or not, smoothed or not", { timeout: 60_000 }, async () => { + const { tf } = await loadTf(); + const random = mulberry32(3), n = 64; + const z = Float64Array.from({ length: n }, () => 6 * random() - 3); + const y = Float32Array.from({ length: n }, () => (random() < 0.3 ? 1 : 0)); + const m = Float32Array.from({ length: n }, () => (random() < 0.7 ? 1 : 0)); + for (const lossType of ["bce", "mse", "margin"] as const) for (const balanceClasses of [false, true]) for (const labelSmoothing of [0, 0.1]) { + const o = { lossType, margin: 0.3, balanceClasses, labelSmoothing, distillationLambda: 0, distillationScope: "unlabeled" as const }; + const expected = batchLoss(z, y, m, o, new Float64Array(n)); + const got = tf.tidy(() => supervisedLoss(tf, tf.tensor1d(Float32Array.from(z)), tf.tensor1d(y), tf.tensor1d(m), o).dataSync()[0]); + expect(got).toBeCloseTo(expected, 5); + } + }); + + it("holds the logits to the starting network's where the scope says", { timeout: 60_000 }, async () => { + const { tf } = await loadTf(); + const z = [1, 2, 3, 4], z0 = [1, 0, 3, 0], m = [1, 1, 0, 0]; + const term = (scope: "unlabeled" | "all") => tf.tidy(() => distillationLoss(tf, tf.tensor1d(z), tf.tensor1d(z0), distillationWeights(tf, tf.tensor1d(m), scope)).dataSync()[0]); + expect(term("unlabeled")).toBeCloseTo(16 / 2); // the two unpainted voxels: 0 and 16 + expect(term("all")).toBeCloseTo((4 + 16) / 4); + }); +}); + +describe("the whole network's training controls", () => { + it("reads the head's loss fields and augmentation into the run's options and the job's params", () => { + const read = readTraining({ finetune_mode: "full", num_epochs: 4, loss_type: "mse", label_smoothing: 0.1, balance_classes: false, distillation_lambda: 0.1, distillation_scope: "all", augment: true, batch_size: 2 }); + if ("error" in read) throw new Error(read.error); + expect(read.options).toMatchObject({ mode: "full", lossType: "mse", labelSmoothing: 0.1, balanceClasses: false, distillationLambda: 0.1, distillationScope: "all", augment: true, batchPatches: 2 }); + expect(read.params).toMatchObject({ loss_type: "mse", label_smoothing: 0.1, balance_classes: false, distillation_lambda: 0.1, distillation_scope: "all", augment: true, finetune_mode: "full" }); + const head = readTraining({ augment: true }); + if ("error" in head) throw new Error(head.error); + expect(head.options.augment).toBe(false); // the head trains on features computed once + }); +}); diff --git a/web/test/unit/finetune-full.test.ts b/web/test/unit/finetune-full.test.ts index 27a42f9e..22a2e52e 100644 --- a/web/test/unit/finetune-full.test.ts +++ b/web/test/unit/finetune-full.test.ts @@ -16,6 +16,7 @@ import { describe, expect, it } from "vitest"; import type { Tf } from "../../src/finetune/full/cellmap-tfjs-train"; import { buildNet, maskedBce, readOnnxWeights, trainStep, writeOnnxWeights } from "../../src/finetune/full/cellmap-tfjs-train"; import { parseTrainGraph, tileSides } from "../../src/finetune/full/graph"; +import { fullTrainStep } from "../../src/finetune/full/loss"; import { checkOffsets, onnxInitializerOffsets } from "../../src/finetune/full/run-full"; import { loadTf } from "../../src/finetune/full/tfjs"; import { readFixture } from "./models-helpers"; @@ -94,6 +95,33 @@ describe("Marwan's TF.js runtime on head.onnx", () => { expect(losses.at(-1)!).toBeLessThan(losses[0]); }); + it("holds the network to its starting self: the distillation term is 0 before the first step and grows as it trains", { timeout: 120_000 }, async () => { + const { tf } = await loadTf(); + const bytes = new Uint8Array(readFileSync(resolve(FIXTURES, "head.onnx"))); + const offsets = onnxInitializerOffsets(bytes), names = Object.keys(TINY.params); + const net = buildNet(tf, TINY, readOnnxWeights(bytes, offsets, names)); + const teacher = buildNet(tf, TINY, readOnnxWeights(bytes, offsets, names), { trainable: new Set() }); + const x = tf.tensor4d(slice(8), [1, 8, 8, 1]); + const target = tf.tensor4d(Float32Array.from({ length: 72 }, (_, i) => (i % 3 === 0 ? 1 : 0)), [1, 6, 6, 2]); + const mask = tf.tensor4d(Float32Array.from({ length: 72 }, (_, i) => (i % 5 === 0 ? 0 : 1)), [1, 6, 6, 2]); + const z0 = tf.tidy(() => teacher.forward(x)); + const bce = tf.tidy(() => maskedBce(tf, z0, target, mask).dataSync()[0]); + const opt = tf.train.adam(0.05); + const o = { lossType: "bce" as const, margin: 0.3, labelSmoothing: 0, balanceClasses: false, distillationLambda: 1, distillationScope: "unlabeled" as const }; + const terms = [0, 1, 2].map(() => { + const step = fullTrainStep(tf, net, opt, x, target, mask, z0, o); + const out = { supervised: step.supervised.dataSync()[0], distillation: step.distillation!.dataSync()[0], total: step.total.dataSync()[0] }; + [step.total, step.supervised, step.distillation!].forEach((s) => s.dispose()); + return out; + }); + expect(terms[0].distillation).toBeCloseTo(0, 6); + expect(terms[0].supervised).toBeCloseTo(bce, 5); // without the head's options, Marwan's masked BCE + expect(terms[2].distillation).toBeGreaterThan(0); + expect(terms[2].total).toBeCloseTo(terms[2].supervised + terms[2].distillation, 5); + [x, target, mask, z0].forEach((t) => t.dispose()); + opt.dispose(); net.dispose(); teacher.dispose(); + }); + it("writes the trained weights back where onnxruntime reads them, in fp32 and fp16", { timeout: 120_000 }, async () => { const { tf } = await loadTf(); for (const [file, tolerance] of [["head.onnx", 1e-5], ["head_fp16.onnx", 5e-3]] as const) {