Exporting to ONNX¶

Trained SkiNet checkpoints can be exported to ONNX format for deployment on mobile (iOS/Android) or any ONNX-compatible inference runtime. The entry point is export_onnx.py in the repo root.

Quick start¶

From an MLflow run folder (recommended)¶

Point --run at the MLflow run directory and the script auto-discovers the best checkpoint and the config YAML:

cd repos/SkiNet
python export_onnx.py \
    --run mlruns/<experiment>/<run_id>/<uuid> \
    --out skinet_unet.onnx

Explicit paths¶

Supply --ckpt and --config directly when you want to target a specific checkpoint or config that is not under artifacts/:

python export_onnx.py \
    --ckpt path/to/epoch=42.ckpt \
    --config path/to/config.yaml \
    --out skinet_unet.onnx \
    --opset 17

CLI reference¶

Flag	Default	Description
`--run`	`None`	MLflow run folder; checkpoint and config are auto-discovered
`--ckpt`	`None`	Explicit path to a `.ckpt` file (required if `--run` is omitted)
`--config`	`None`	Explicit path to a config YAML (required if `--run` is omitted)
`--out`	`skinet_unet.onnx`	Output path for the exported ONNX model
`--opset`	`17`	ONNX opset version

Either --run or both --ckpt and --config must be supplied. If --run is given alongside --ckpt/--config, --run takes precedence and the explicit paths are silently ignored.

What the script does¶

Loads the config via SkiNet.ML.configs.load_config_from_yaml.load_config_from_yaml().
Builds the model architecture from the config with use_torch_compile forced to False (compilation is not needed for export and avoids environment dependencies).
Loads the checkpoint with torch.load(..., weights_only=False) and strips the model._orig_mod.* key prefix that torch.compile adds, so the state dict matches the uncompiled model.
Reports the optimal threshold stored in the checkpoint buffer (optimal_threshold). This value should be hard-coded as SEGMENTATION_THRESHOLD in the iOS/Android app.
Wraps the backbone in _UNetWithSigmoid, which fuses a torch.sigmoid into the ONNX graph so the model outputs probabilities (0–1) rather than raw logits.
Exports via torch.onnx.export with dynamic batch axis on both input and output.
Merges external weight data (if onnx is installed): the dynamo exporter may write a .onnx.data sidecar file; the script merges it into a single self-contained .onnx file and deletes the sidecar.
Validates with ONNXRuntime (if onnxruntime is installed): runs a zero-tensor forward pass and asserts the output shape is (1, 1, 256, 256).
Prints a deployment summary: model file size, INPUT_SIZE, normalisation constants, and the optimal sigmoid threshold.

ONNX graph¶

Property	Value
Input name	`image`
Input shape	`(batch, 3, 256, 256)` — only the batch axis is dynamic
Output name	`mask_prob`
Output shape	`(batch, 1, 256, 256)` — probabilities in [0, 1]
Default opset	`17`

Note

Only the batch axis is dynamic. The spatial dimensions are fixed at 256×256, so every input must be resized to 256×256 before inference — feeding any other height/width will fail ONNXRuntime’s shape check. 256 is the size the model is trained and deployed at.

iOS / mobile preprocessing constants¶

The script prints these at the end of every run:

INPUT_SIZE  = 256
NORM_MEAN   = [0.699, 0.556, 0.5121]
NORM_STD    = [0.1576, 0.1562, 0.1706]
THRESHOLD   = <value from checkpoint>

Apply these in the same order as the training pipeline:

Resize the input image to INPUT_SIZE × INPUT_SIZE.
Normalise each channel: (pixel / 255 − mean) / std.
Run inference; apply the threshold to the probability map to get a binary mask.

Auto-discovery rules (`--run` mode)¶

When --run is supplied, _resolve_run searches the run folder as follows:

Checkpoint: glob("artifacts/checkpoints/**/*.ckpt"), sorted lexicographically. If any result has "best" in its path components, the last such path is chosen; otherwise the lexicographically last checkpoint overall is used.
Config: glob("artifacts/config/*.yaml"), sorted lexicographically; the last file is used.

Optional dependencies¶

Package	Effect if missing
`onnx`	Weight-merge step is skipped; model may export as two files (`.onnx` + `.onnx.data`)
`onnxruntime`	Post-export shape validation is skipped

Install both for a full export pipeline:

pip install onnx onnxruntime

Checkpoint key remapping¶

Training may be run with use_torch_compile: true, which causes torch.compile to prefix every state-dict key with model._orig_mod.. The exporter strips this prefix automatically, so compiled and uncompiled checkpoints are both supported without any manual key editing.