MUTANT

Multi-objective UNet Tiny Architecture NAS Tool

MUTANT searches compact U-Net segmentation architectures with an evolutionary, multi-objective neural architecture search, then trains the winners from scratch and quantizes them to INT8 for edge deployment. It is dataset-agnostic: describe your tiles in a dataset.json, choose the number of classes and the class whose IoU matters, and search.

• Search (mutant search): genetic algorithm over a token grammar of U-Net blocks (standard / depthwise-separable / residual convolutions, dilations, GELU, GroupNorm, searchable squeeze-excitation, ASPP and attention-gate genes). Candidates are pre-filtered with zero-cost proxies, trained briefly under Successive Halving, scored on the validation IoU of the target class, and ranked with NSGA-II over {IoU ↑, FLOPs ↓} with the parameter budget enforced as constrained dominance. Output: a Pareto front of genotypes.
• Train (mutant train): full training of a genotype (FP32 with early stopping → temperature + threshold calibration → quantization-aware training with the PyTorch 2 export flow of torchao → INT8 model → test metrics, CPU latency and weight size).
• Evaluate (mutant eval): FP32 checkpoints and INT8 programs on any split.

In the case study that produced the tool, a two-class segmentation task on a public multispectral satellite dataset, a 626 k-parameter finalist reached a target-class test IoU of 0.804 in INT8, within ±0.005 of a 7 M-parameter reference U-Net, and a 1 182-parameter network still reached 0.79 (single seed, see docs/06_case_study.md for the numbers and their caveats).

Install

python -m venv .venv && source .venv/bin/activate
pip install torch --index-url https://download.pytorch.org/whl/cpu   # or a CUDA wheel
pip install -e ".[dev]"          # distribution name: mutant-nas

Extras: [int8] (torch ≥ 2.11 + torchao, quantization-aware training), [geotiff] (rasterio, scene-level datasets), [mlflow], [viz] (figures), [dev] (tests, lint, everything above). import mutant and mutant --help work without torch installed.

Quick start

mutant --help
mutant search --smoke                    # end-to-end search on the built-in synthetic dataset (CPU, seconds)
mutant data synth ./demo --in-channels 4 --num-classes 3 --tile-size 64   # a synthetic dataset on disk
mutant data check ./demo                 # validate dataset.json, splits, shapes, class histogram
mutant data stats ./demo                 # per-channel mean/std → dataset.json
mutant search --data ./demo --smoke      # search on a dataset directory
mutant search --help                     # all search options (population, generations, budget, ...)
mutant train --data ./demo --genotype "B8 E[+SE Cd3g Cd3rs] BN[Rc3r] D[Cd3rm]" --smoke   # FP32 → QAT → INT8 → test
mutant eval  --data ./demo --genotype "B8 E[+SE Cd3g Cd3rs] BN[Rc3r] D[Cd3rm]" \
             --checkpoint runs/train_finalist/best_fp32.pth --int8 runs/train_finalist/int8_model.pt2

Outputs (runs/*.json, runs/train_<name>/, optional mlflow/) are written under the working directory (override with MUTANT_OUTPUT_DIR). Drop --smoke for real runs.

The same steps, with explanations and their outputs, are in notebooks/00_quickstart.ipynb (synthetic data, CPU, a few minutes): dataset → genotype → search → training with INT8 → evaluation → how to plug in your own data.

Bring your own dataset

A dataset is a directory with a dataset.json (channels, classes, target class, tile size, normalization) and one of two layouts:

• tiles: <split>/*.npz files with image [C,H,W] and mask [H,W] (see examples/convert_to_tiles.py);
• geotiff: <split>/images/*.tif scenes and <split>/masks/*.tif label rasters, tiled on the fly (pip install -e ".[geotiff]", i.e. rasterio); labels are decoded with a label_map or a custom label_fn declared in dataset.json.

Mask values are class indexes 0..num_classes-1, with -1 for pixels to ignore. Details in docs/04_data.md.

Documentation

Document	Content
01_overview.md	Problem, the two phases, design principles, positioning, limitations
02_search_space.md	Token grammar, genes, genotype strings, how a genotype becomes a network
03_algorithm.md	Pre-filter, Successive Halving, fitness, NSGA-II, reproduction, run schema
04_data.md	dataset.json, tiles and geotiff adapters, label decoding, mutant data
05_training_and_int8.md	Full training protocol, calibration, QAT with torchao, INT8 artifacts, mutant eval
06_case_study.md	The case study: searches, finalists, reference model, caveats, how to reproduce
07_reproducibility.md	Install matrix, seeds, outputs, MLflow, CI, timings
08_api.md	Python API with runnable examples
PROVENANCE.md	Where the code comes from and what was verified
ROADMAP.md	What is not there yet

Notebooks: notebooks/00_quickstart.ipynb (executable walk-through of the tool on synthetic data), notebooks/01_search_analysis.ipynb (figures of the search runs) and notebooks/02_case_study_results.ipynb (finalists). Figures in docs/figures/ are generated by scripts/build_figures.py from runs/case_study/.

Design in brief

Choice	Why
Two separate phases (search → full training)	Short trainings rank architectures; only full training produces results. Mixing the two couples variables that must be analysed separately.
Zero-cost proxies only as a conservative pre-filter (≤ 30 %)	On dense prediction their correlation with accuracy is weak, especially among the best candidates (Krishnakumar et al. 2022).
Fixed search loss (Dice + Focal)	The loss depends more on the data than on the architecture; co-searching it adds noise to the ranking.
Fitness from a global confusion matrix	Per-image IoU averages are biased (PASCAL VOC protocol, Everingham et al. 2010).
NSGA-II with the parameter budget as constrained dominance	Rejecting oversized candidates throws away gradient information near the boundary; dominance keeps the pressure towards small models (Deb 2000).
Aging + hybrid elitism + variable-length crossover	Regularization against lucky individuals in a noisy few-epoch regime (Real et al. 2019); valid children between parents of different depth.
Per-model temperature + threshold calibration	A monotone rescaling that makes the operating point transferable from validation to test and from FP32 to INT8 (Guo et al. 2017).
INT8 through the PyTorch 2 export flow	The successor of FX graph-mode quantization; the converted model is persisted as a .pt2 program because activation quantization parameters are graph constants.

Status and limitations

• Search, data adapters, train and eval are functional and tested end to end (CPU tests in CI; a full search needs a GPU and hours).
• Single-GPU, train-from-scratch regime: no weight sharing, no supernet. Short-training fitness is noisy; results must be confirmed by full training of the finalists.
• CNN-only vocabulary; attention tokens and a frozen-encoder profile are in the roadmap.
• The case study is single-seed and has no random-search control; read its caveats.

Development

ruff check src tests examples && pytest -q && pytest -q -m slow && bash scripts/release_check.sh

See CONTRIBUTING.md and CHANGELOG.md.

Citation

See CITATION.cff.

License

MIT.