cellmap_flow.finetune.options
Every setting a finetune trains with, and its default, in one place.
TrainingOptions is what the Finetune tab’s form sends
(dashboard/requests.FinetuneSubmit adds the request’s own fields to
it), what the job manager turns into the job’s command line
(cli.training_flags), what that command line parses back into
(cli.build_arg_parser makes a flag of every field, cli.options_of
reads them), and what the trainer is built from
(trainer.Finetuner). Each default is written here once:
cellmap_flow.schema exports the form’s to the browser as
FINETUNE_DEFAULTS, so the browser’s form starts from the same values.
The settings are kept as they are asked for. for_session says what a
job can train with given its session (sparse annotations, a distance or a
flow target); the trainer applies the rest of its own rules (label
smoothing means nothing to the margin loss, say) when it is built.
Attributes
Classes
The settings of one finetune, each with its default. |
Module Contents
- cellmap_flow.finetune.options.LOSS_TYPES = ('dice', 'bce', 'combined', 'mse', 'margin', 'interval', 'flow')
- cellmap_flow.finetune.options.OUTPUT_TYPES = ('binary', 'binary_broadcast', 'affinities', 'distance', 'flows')
- cellmap_flow.finetune.options.MODES = ('lora', 'full', 'head')
- cellmap_flow.finetune.options.OUTPUT_CLASSES
- cellmap_flow.finetune.options.MANIFEST_FIELDS = ('patches_per_epoch', 'rehearsal_fraction')
- cellmap_flow.finetune.options.LEARNING_RATE_FLOOR = 0.1
- class cellmap_flow.finetune.options.TrainingOptions(/, **data: Any)
The settings of one finetune, each with its default.
The defaults are the server form’s (
_finetune_tab.html, checked by tests/web/test_schema.py): a LoRA adapter of rank 64, 20 epochs of batches of 8 from 1e-4 eased down to a tenth of it, the margin loss with class balance and a light distillation (0.01) on the unlabeled voxels, no augmentation. The browser’s form overrides a few for the last layer (web/src/finetune/types.ts DEFAULT_OPTIONS).- model_config
Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].
- mode: Mode
- lora_r: int
- lora_alpha: int | None
- lora_dropout: float
- lora_min_channels: int
- num_epochs: int
- batch_size: int
- learning_rate: float
- learning_rate_schedule: LearningRateSchedule
- gradient_accumulation_steps: int
- loss_type: LossType
- label_smoothing: float
- distillation_lambda: float | None
- distillation_scope: DistillationScope
- margin: float
- balance_classes: bool
- slope_weight: float
- augment: bool
- mixed_precision: bool
- num_workers: int
- tensorboard: bool
- mask_unannotated: bool
- output_type: OutputType
- distance_sigma: float
- select_channel: int | None
- offsets: str | None
- output_class: OutputClass | None
- patches_per_epoch: int | None
- rehearsal_fraction: float | None
- replaced(**changes) TrainingOptions
A copy with
changesapplied, validated as a whole.
- for_session(*, sparse: bool) TrainingOptions
What a job trains with, given whether its session is sparse (painted).
Some combinations cannot train, or would train the wrong thing, so they are replaced here, where the dashboard shows it in submit’s answer, rather than failing on the cluster: - sparse annotations with mse: the margin loss, with distillation to
the base model at 0.5, which is how sparse annotations train (see the distance case below);
a distance target with sparse annotations: the interval loss, on the bounds the paint sets on each voxel’s distance, with distillation of at least 0.5 (below);
a distance target otherwise: bce, without label smoothing;
a flow target (Cellpose): the flow loss, the only one it goes with, without label smoothing, sparse or not (it masks per channel itself).
A sparse session also masks its unannotated voxels out of the loss.