cellmap_flow.finetune.distill ============================= .. py:module:: cellmap_flow.finetune.distill .. autoapi-nested-parse:: 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 ````: - ``-small/manifest.json`` and its ONNX files, as export_web writes a model's browser form (```` 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 ``/web_export/-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 \ --corrections /corrections --output \ [--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 ---------- .. autoapisummary:: cellmap_flow.finetune.distill.logger cellmap_flow.finetune.distill.MODULE cellmap_flow.finetune.distill.CHECKPOINT_NAME cellmap_flow.finetune.distill.BROWSER_PARAMETERS cellmap_flow.finetune.distill.TEACHER_SHARE cellmap_flow.finetune.distill.RANDOM_SHARE cellmap_flow.finetune.distill.HELD_OUT_BLOCKS cellmap_flow.finetune.distill.HELD_OUT_SEED cellmap_flow.finetune.distill.EPOCHS cellmap_flow.finetune.distill.HOURS cellmap_flow.finetune.distill.STEPS_PER_EPOCH cellmap_flow.finetune.distill.LEARNING_RATE cellmap_flow.finetune.distill.WORKERS cellmap_flow.finetune.distill.BROWSER_SIDES cellmap_flow.finetune.distill.SETUP_HOURS Functions --------- .. autoapisummary:: cellmap_flow.finetune.distill.make_student cellmap_flow.finetune.distill.closeness cellmap_flow.finetune.distill.distill cellmap_flow.finetune.distill.start_distill cellmap_flow.finetune.distill.main Module Contents --------------- .. py:data:: logger .. py:data:: MODULE :value: 'cellmap_flow.finetune.distill' .. py:data:: CHECKPOINT_NAME :value: 'student_best.pt' .. py:data:: BROWSER_PARAMETERS :value: 12000000 .. py:data:: TEACHER_SHARE :value: 0.1 .. py:data:: RANDOM_SHARE :value: 0.5 .. py:data:: HELD_OUT_BLOCKS :value: 8 .. py:data:: HELD_OUT_SEED :value: 1000000007 .. py:data:: EPOCHS .. py:data:: HOURS .. py:data:: STEPS_PER_EPOCH :value: 100 .. py:data:: LEARNING_RATE :value: 0.0001 .. py:data:: WORKERS :value: 4 .. py:data:: BROWSER_SIDES .. py:data:: SETUP_HOURS :value: 1.0 .. py:function:: 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. .. py:function:: 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. .. py:function:: 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. .. py:function:: 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 ``/web_export/-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. :raises RuntimeError: the job could not be submitted. .. py:function:: main(argv=None) -> int