Whole-network finetuning in the browser: the last layer's loss settings, and augmentation

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=<url>&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.

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

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

cellmap_flow/dashboard/templates/_finetune_tab.html

@@ -541,9 +541,10 @@ <h5>Training Configuration</h5>
        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 @@ <h5>Training Configuration</h5>
       </div>
     </div>
 
-    {# The whole network trains with masked BCE on the paint alone: the rest are the head's. #}
-    <div class="row head-only">
+    <div class="row">
       <div class="col-md-6">
         <div class="mb-3">
           <label for="lossType" class="form-label">Loss Function</label>
@@ -766,7 +766,7 @@ <h5>Training Configuration</h5>
       <div class="col-md-6">
         <div class="mb-3">
           {% if mode == "browser" %}
-          <label for="distillationLambda" class="form-label" title="Adds this weight times the mean squared distance of the head's logits from the starting head's, on the voxels the scope below picks (or, once good regions are marked, their unpainted voxels): what the paint does not ask to change stays as the model had it. 0.01 is a light hold; a head moves freely, so 0.1 to 1 holds it firmly.">
+          <label for="distillationLambda" class="form-label" title="Adds this weight times the mean squared distance of the model's logits from the starting model's, on the voxels the scope below picks (for the last layer, once good regions are marked, their unpainted voxels): what the paint does not ask to change stays as the model had it. 0.01 is a light hold; a head moves freely, so 0.1 to 1 holds it firmly. The whole network is held on the same tile it trains on, as the server holds it.">
             Distillation Weight
           </label>
           {% else %}
@@ -787,14 +787,14 @@ <h5>Training Configuration</h5>
             <option value="10.0">10.0</option>
           </select>
           {% if mode == "browser" %}
-          <small class="text-muted">Keeps the head near the model's own output where nothing is painted.</small>
+          <small class="text-muted">Keeps the model near its own output where nothing is painted.</small>
           {% else %}
           <small class="text-muted">Keeps the model near its original output. Keep it small.</small>
           {% endif %}
         </div>
       </div>
     </div>
-    <div class="row head-only">
+    <div class="row">
       <div class="col-md-6">
         <div class="mb-3">
           <label for="distillationScope" class="form-label">
@@ -837,12 +837,19 @@ <h5>Training Configuration</h5>
           </div>
           <small class="text-muted">Weights fg and bg equally.</small>
         </div>
-        <div class="mb-3 server-only">
+        {# 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. #}
+        <div class="mb-3 server-only full-also">
           <div class="form-check">
             <input class="form-check-input" type="checkbox" id="augment">
+            {% if mode == "browser" %}
+            <label class="form-check-label" for="augment" title="Each step's tiles get a random brightness (×0.8 to ×1.2) and noise (1% of the range), then random flips and quarter turns of the image and your paint together, with the targets made after, as on the server. The whole network sees the same tiles every epoch, so without it, it can learn them by heart. Off by default, as on the server.">
+              Augmentation
+            </label>
+            {% else %}
             <label class="form-check-label" for="augment" title="Random flips, XY rotations, brightness and noise, on top of the patch-centre jitter that is always applied. Off by default: augmentation earns its keep when a run sees the same patches many times over, so a short correction run mostly just gets noisier gradients.">
               Augmentation
             </label>
+            {% endif %}
           </div>
           <small class="text-muted">
             <span id="augmentAdvice" class="d-block"></span>

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

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