cellmap_flow.finetune.distill

Distill a finetuned model into a small copy the browser can run, and export the copy for the browser.

The large cellmap UNets (about 800M parameters, 3.2 GB, such as mito-aff-unet-setup-16) are too big for a browser to load, and export_web gives them no browser form. This trains a small network to compute what a finetuned model’s network computes, and exports it as export_web exports a model’s browser form. The network surgery is cellmap-models’ (its distill_student); this is the job around it:

The teacher is the network the serving YAML serves, as export_web builds it (export_web.finetuned_network): the base as the trainer builds it, the finetune folded in, every layer a real one.
The student is cellmap-models’ distill_student: for a funlib UNet (what prune_funlib_unet covers, such as the cellmap StandardUnets) the same network with fewer channels at every level, started from the teacher’s strongest channels (init “pruned”); for any other network (ResUNet, Unet3dIso…) a fly_organelles StandardUnet with the teacher’s channels and output activation, started at random (init “random”), at the largest block whose output fits inside the teacher’s. A random start has more to learn, so it gets more epochs and hours by default.
The distillation trains the student to reproduce the teacher’s output as the network emits it, the mean squared difference over every voxel and channel (losses.distillation_loss, the finetune’s teacher term), on blocks of the raw the model was finetuned on, read and normalized as the trainer reads its patches (data.loader.dataset_from_manifest): half centred on what was annotated, half at random points of the volume. It needs no annotations. Each block is read large enough for both networks’ inputs; both read and write around the same centre, as the inference runner places a model’s output in the middle of what it reads (ModelGeometry.context), so the teacher’s input and the student’s are the middle of the block, and the student’s target is the middle of the teacher’s output. Squared difference of the raw output because the output may be logits, affinities or distances, which the network cannot tell apart (probe_output_class calls all three unbounded): what is used downstream of affinities and distances is their value, every level of it (watershed and agglomeration thresholds, distance cut-offs), and a student that matches the values matches any threshold or postprocessing on them, logits’ included. A sigmoid’s cross-entropy would read distances as logits and stop caring where both are far from 0. Adam takes the loss’s scale out of its steps, so logits of about 5 or about 1400 (export_merged’s finetunes) train alike.
It runs on the GPU when there is one (utils.device.best_device), the training steps in mixed precision (trainer.recovery.Precision: bf16 where the GPU has it), the held-out comparison in float32 with TF32 off, so that its numbers are the networks’ own.
How close: on held-out blocks drawn apart from the training’s, cellmap-models’ output_difference of the student from the teacher (the export’s metrics: the mean squared, mean and largest absolute difference, and the share of voxels on the other side of the model’s decision boundary, 0.5 for a model that ends in a sigmoid and else 0, as portable.label_it.probability reads it), for the start and for the student exported.

What it writes in <output>:

<folder>-small/manifest.json and its ONNX files, as export_web writes a model’s browser form (<folder> is export_web’s folder for the finetuned model), with a distillation block: the teacher’s folder (so that a publish can say “distilled from” its ref), the teacher’s parameters, the student’s init, num_fmaps and fmap_inc_factor, the steps it trained, and the held-out closeness of the start and of the student (distilled);
student_best.pt: the best student so far, written each time the held-out difference improves, so a job killed before it exports keeps it; a run that makes the same student continues from it;
result.json and export_log.txt, as export_web’s: {"folder", "manifest"} or {"error"}, also printed as the output’s last line.

The progress lines are the trainer’s, “Starting epoch N of M” and “Epoch N/M - Loss: …” (markers’ patterns), each epoch’s with the held-out closeness after it.

The dashboard starts it and follows it as an export, with export_web’s export_status:

run = start_distill(yaml_path, corrections_dir, run_dir)
export_web.export_status(run)   # the answer start_export's gets

start_distill submits the job as start_export does (export_web.start_job: in the environment the model finetunes in, on its YAML’s queue), into <run_dir>/web_export/<YAML's name>-small/, with the training’s time limit (the longer, a random start’s, unless given) and an hour for the rest as its walltime, and returns at once. From the command line:

python -m cellmap_flow.finetune.distill --yaml <serving yaml> \
    --corrections <session>/corrections --output <dir> \
    [--parameters 12000000] [--epochs 100] [--steps-per-epoch 100] [--hours 3]

The student’s size: at most --parameters, by default a tenth of the teacher’s and no more than BROWSER_PARAMETERS, about what the browser catalog’s largest models have (distill_student says how the widths are chosen). For mito-aff-unet-setup-16 (792M parameters) that is num_fmaps 13 and fmap_inc_factor 3, 10.8M parameters, 43 MB in fp32.

A 2D model (a block one z-slice thick) gets a 2D student (distill_student’s 1x3x3 kernels and 1x2x2 pooling), trained and exported on slices. The browser copy is exported at the block the student trained at and, for a StandardUnet student, at the larger blocks it computes too (BROWSER_SIDES, cellmap-models’ Student.blocks), each checked against the student on real raw.

What it refuses, before it trains: a model that writes at another voxel size than it reads (the student writes at the voxel size it reads); one that takes more than one input channel (the browser feeds one); a teacher served other than by its network’s forward (Cellpose); and a teacher whose output block holds no StandardUnet block.

Needs cellmap-models with model_export.student and model_export.onnx_web (the pin in pixi.toml’s finetune feature) and its export extra, onnx and onnxruntime, where it runs.

Attributes

`logger`
`MODULE`
`CHECKPOINT_NAME`
`BROWSER_PARAMETERS`
`TEACHER_SHARE`
`RANDOM_SHARE`
`HELD_OUT_BLOCKS`
`HELD_OUT_SEED`
`EPOCHS`
`HOURS`
`STEPS_PER_EPOCH`
`LEARNING_RATE`
`WORKERS`
`BROWSER_SIDES`
`SETUP_HOURS`

Functions

`make_student`(teacher, geometry, out_channels, output_class)	cellmap-models' student of `teacher` (its `Student`), for the
`closeness`(network, blocks, targets, input_size, threshold) → dict	How far `network`'s output on `blocks` (the middle `input_size`
`distill`(yaml_path, corrections_dir, output_dir, *[, ...]) → dict	Distill the model `yaml_path` serves into a small copy and export
`start_distill`(yaml_path, corrections_dir, run_dir, *) → cellmap_flow.finetune.export_web.WebExport	Start distilling the model `yaml_path` serves, trained on the
`main`([argv]) → int

Module Contents

cellmap_flow.finetune.distill.logger

cellmap_flow.finetune.distill.MODULE = 'cellmap_flow.finetune.distill'

cellmap_flow.finetune.distill.CHECKPOINT_NAME = 'student_best.pt'

cellmap_flow.finetune.distill.BROWSER_PARAMETERS = 12000000

cellmap_flow.finetune.distill.TEACHER_SHARE = 0.1

cellmap_flow.finetune.distill.RANDOM_SHARE = 0.5

cellmap_flow.finetune.distill.HELD_OUT_BLOCKS = 8

cellmap_flow.finetune.distill.HELD_OUT_SEED = 1000000007

cellmap_flow.finetune.distill.EPOCHS

cellmap_flow.finetune.distill.HOURS

cellmap_flow.finetune.distill.STEPS_PER_EPOCH = 100

cellmap_flow.finetune.distill.LEARNING_RATE = 0.0001

cellmap_flow.finetune.distill.WORKERS = 4

cellmap_flow.finetune.distill.BROWSER_SIDES

cellmap_flow.finetune.distill.SETUP_HOURS = 1.0

cellmap_flow.finetune.distill.make_student(teacher, geometry, out_channels, output_class, parameters=None)

cellmap-models’ student of teacher (its Student), for the model of geometry, whose output has out_channels channels and is of output_class (probe_output_class): at most parameters parameters, by default a tenth of the teacher’s and at most BROWSER_PARAMETERS.

Raises:: ValueError – no student is that small, or one that large would not fit an .onnx file, or no StandardUnet block fits in the teacher’s.

cellmap_flow.finetune.distill.closeness(network, blocks, targets, input_size, threshold: float) → dict: How far network’s output on blocks (the middle input_size of each) is from targets, the teacher’s: cellmap-models’ output_difference (mse, mean_abs, max_abs, flip_frac on threshold) over all of them, computed in float32 with TF32 off.

cellmap_flow.finetune.distill.distill(yaml_path, corrections_dir, output_dir, *, parameters=None, epochs=None, steps_per_epoch=STEPS_PER_EPOCH, hours=None, learning_rate=LEARNING_RATE, workers=WORKERS, variants=None) → dict

Distill the model yaml_path serves into a small copy and export it into output_dir; {"folder", "manifest"} (see the module’s docstring).

Trains for epochs of steps_per_epoch blocks, or until hours have passed, whichever is first (by default EPOCHS and HOURS of how the student starts), and exports the student that came closest to the teacher on the held-out blocks.

Raises:: ValueError – the model is one this cannot distill or export (see the module’s docstring), or its session has no training manifest.

cellmap_flow.finetune.distill.start_distill(yaml_path, corrections_dir, run_dir, *, parameters=None, hours=None) → cellmap_flow.finetune.export_web.WebExport

Start distilling the model yaml_path serves, trained on the session corrections_dir, into <run_dir>/web_export/<YAML's name>-small/; follow it with export_web.export_status.

parameters: the student’s largest size (by default make_student’s); hours: the training’s time limit (by default HOURS of how the student starts, which the job decides; the walltime allows the longer). The job runs as start_export’s does (export_web.start_job), in the environment the model finetunes in on the YAML’s queue and charge group, and a result an earlier run left there is removed first, so that export_status does not read it as this one’s; its best student is kept, for this run to continue from.

Raises:

ValueError – the YAML serves no finetuned model, or its environment cannot run the job.
RuntimeError – the job could not be submitted.

cellmap_flow.finetune.distill.main(argv=None) → int