ES-MoE Toolkit

_{Docs: English · 中文 · Quick start in Colab: notebooks/quickstart.ipynb · Ask questions about the code: DeepWiki.}

Install

pip install esmoe

The distribution, the import and the CLI are all esmoe; the project is written ES-MoE in prose.

Requires ultralytics.

Use

One call covers register, graft, build and wire:

import esmoe

model = esmoe.equip("yolo11n.yaml", weight=0.01)
model.train(data="coco8.yaml", epochs=10)

Or take the steps apart when you need control over each one:

from ultralytics import YOLO
import esmoe

esmoe.inject_esmoe()                                    # make `ESMoE` resolvable in model.yaml
esmoe.graft("yolov8n.yaml", out="v8-esmoe.yaml", at=[4, 6])  # insert blocks, renumber the head
model = YOLO("v8-esmoe.yaml")
esmoe.attach_aux_loss(model, weight=0.01)               # router loss joins the training loss

From the shell:

esmoe graft yolo11n.yaml -o yolo11n-esmoe.yaml -e 4 -k 2 --at backbone_end

attach_aux_loss adds an esmoe_aux entry to the trainer's loss table, so a non-zero, back-propagated auxiliary term shows up in results.csv rather than merely in a config.

Written by hand, a grafted config layer is just:

[-1, 1, ESMoE, [4, 2]]   # num_experts, top_k

The block is channel preserving and infers its width on the first forward, which is what lets stock parse_model size it without a patch.

Extend

Experts and the balancing objective are plain callables, so a variant is a few lines:

esmoe.ESMoE(num_experts=4, top_k=2, expert=MyExpert, balance=my_balance_fn)

MyExpert(c1, c2, k) -> Module, my_balance_fn(probs, gate) -> scalar. esmoe.blocks(model) walks every block in a model, and esmoe.collect_aux_loss(model) returns the current step's router loss for custom training loops.

Compatibility

Verified by tests/test_ultralytics.py on ultralytics 8.4.101 and 8.4.132, which report loss items in two different shapes; both are handled. The training column is backed by real 1-epoch VisDrone runs on all four generations (results/*-compat-*.json), each logging a non-zero train/esmoe_aux.

DDP works: attach_aux_loss routes model.train() through a trainer class that lives in esmoe.trainer, so the worker processes ultralytics spawns register the block and the auxiliary loss before they build. Verified by scripts/verify.py (the real worker file in a fresh interpreter; two gloo ranks with agreeing router gradients). Not supported together with compile=True, which turns off find_unused_parameters.

Selected default

ESMoE(num_experts=4, top_k=2) with attach_aux_loss(weight=0.01), chosen under one budget over 2/4/8-expert and top-1 variants. Under the repository protocol (VisDrone, imgsz 800, 120 epochs, three seeds) the paired result on YOLOv8n is +0.0025 mAP50 (2/3 seeds) and +0.0004 mAP50-95, at +10.4% parameters and about 9% more wall-clock per epoch. By COCO-style area bucket, large objects get consistently worse (APl −0.010, 0/3) and small-object recall consistently better (ARs +0.003, 3/3). On YOLO11n the same comparison is a wash: the mechanics transfer across backbones, the accuracy effect does not.

The default graft leaves YOLOv8's P5 lateral reading SPPF rather than the block; graft(..., rewire=True) changes that and is the next comparison to run. Reasoning and full tables: docs/SELECTION.md, results/buckets.md, results/routing.md.

Develop and reproduce

uv sync --group dev
uv run pytest -q
uv run python scripts/capture_env.py                  # freeze environment into env/
EPOCHS=20 FRACTION=0.25 SEEDS="0 1 2" bash scripts/sweep.sh
uv run python scripts/report.py                       # results/summary.md

Every run writes one machine-readable record to results/ (config, dataset, hardware, budget, seed, metrics, artifact, status, limitation). Read limitations.md before quoting any number.

Linked projects

• ultralytics - the official YOLO framework this plugs into.
• Tencent/YOLO-Master - where ES-MoE comes from (paper).

License

AGPL-3.0-only, matching the Ultralytics ecosystem it builds on.