Upload folder using huggingface_hub

README.md

@@ -1,13 +1,39 @@
 ---
-title: Mopet Medical Classification
-emoji: 👁
-colorFrom: blue
-colorTo: green
+title: MoPET Medical Classification
+emoji: 🩺
+colorFrom: gray
+colorTo: pink
 sdk: gradio
 sdk_version: 6.22.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: MoPET mixture-of-experts medical image classification
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MoPET Medical Classification
+
+This Space demos **MoPET: Parameter-Efficient Mixture-of-Experts for Unified
+Medical Image Classification** (EMA4MICCAI 2026 Workshop).
+
+MoPET adapts a *frozen* DINOv3 ViT-B/16 backbone with a learned sparse top-k
+router over a heterogeneous pool of LoRA + BOFT PEFT experts injected into the
+attention `qkv` projections. A single model consolidates four MedMNIST+
+classification tasks (Blood, Breast, Derma, Path) behind one shared,
+sparsely-routed expert pool.
+
+## Usage
+
+1. Upload a medical image (blood cell microscopy, breast ultrasound, dermoscopy,
+   or colon pathology histology).
+2. Select the matching task head.
+3. Click **Classify** to get per-class probabilities.
+
+> **Disclaimer:** This is a research artifact, not a diagnostic device. Outputs
+> must not be used for clinical diagnosis.
+
+## Links
+
+- [Paper (arXiv)](https://arxiv.org/abs/2607.29462)
+- [GitHub](https://github.com/sdoerrich97/mopet)
+- [Pretrained Weights](https://huggingface.co/sdoerrich97/mopet_dinov3_unified_blood_breast_derma_path)

app.py

@@ -0,0 +1,188 @@
+"""MoPET: Parameter-Efficient Mixture-of-Experts for Unified Medical Image Classification.
+
+A Gradio demo that loads the pretrained MoPET ``unified`` checkpoint and lets
+visitors upload a medical image, pick one of four MedMNIST+ tasks, and receive
+a classification from the model's per-dataset head.
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # MUST come before torch / any CUDA-touching import
+
+import gradio as gr
+import timm
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torchvision import transforms
+
+from mopet import create_model
+from mopet._factory import BACKBONES, PUBLISHED_MODELS
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+VARIANT = "unified"
+PUBLISHED = PUBLISHED_MODELS[VARIANT]
+DATASET_NAMES = list(PUBLISHED.datasets)  # head order = dataset id
+
+# MedMNIST+ class labels (from the official INFO table)
+CLASS_LABELS: dict[str, list[str]] = {
+    "BloodMNIST": [
+        "basophil",
+        "eosinophil",
+        "erythroblast",
+        "immature granulocytes",
+        "lymphocyte",
+        "monocyte",
+        "neutrophil",
+        "platelet",
+    ],
+    "BreastMNIST": ["malignant", "normal, benign"],
+    "DermaMNIST": [
+        "actinic keratoses",
+        "basal cell carcinoma",
+        "benign keratosis-like lesions",
+        "dermatofibroma",
+        "melanoma",
+        "melanocytic nevi",
+        "vascular lesions",
+    ],
+    "PathMNIST": [
+        "adipose",
+        "background",
+        "debris",
+        "lymphocytes",
+        "mucus",
+        "smooth muscle",
+        "normal colon mucosa",
+        "cancer-associated stroma",
+        "colorectal adenocarcinoma epithelium",
+    ],
+}
+
+# DINOv3 backbone normalization stats (from timm pretrained config)
+_tim_id = BACKBONES[PUBLISHED.backbone]
+_cfg = timm.get_pretrained_cfg(_tim_id)
+_MEAN = tuple(_cfg.mean)
+_STD = tuple(_cfg.std)
+
+# Preprocessing: ToTensor → Normalize → Pad to 224 → Resize to 256
+# (matches the paper's `build_transform` for DINOv3 at resolution 224)
+_padding = max(0, 224 - 224)  # 0 when loading at 224; images are resized anyway
+_pad_l = _pad_t = _padding // 2
+_pad_r = _padding - _pad_l
+_pad_b = _padding - _pad_t
+
+TRANSFORM = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=_MEAN, std=_STD),
+    transforms.Pad((_pad_l, _pad_t, _pad_r, _pad_b), fill=0, padding_mode="constant"),
+    transforms.Resize((256, 256)),
+])
+
+# ---------------------------------------------------------------------------
+# Model loading (module scope, eager .to("cuda"))
+# ---------------------------------------------------------------------------
+
+print(f"Loading MoPET '{VARIANT}' (backbone={PUBLISHED.backbone}, datasets={DATASET_NAMES}) ...")
+model = create_model(weights=VARIANT, map_location="cpu").eval()
+model = model.to("cuda")
+print("Model loaded and moved to CUDA.")
+
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+
+@spaces.GPU(duration=60)
+def classify(image: Image.Image, dataset_name: str) -> dict:
+    """Classify a medical image using the MoPET unified model.
+
+    Args:
+        image: An RGB medical image (blood cell, breast ultrasound, dermoscopy,
+            or colon pathology histology).
+        dataset_name: Which MedMNIST+ task head to use for classification.
+
+    Returns:
+        A label-probability dictionary for the selected task.
+    """
+    if image is None:
+        return {label: 0.0 for label in CLASS_LABELS[dataset_name]}
+
+    image = image.convert("RGB")
+    tensor = TRANSFORM(image).unsqueeze(0).to("cuda")
+
+    dataset_id = DATASET_NAMES.index(dataset_name)
+    dataset_ids = torch.tensor([dataset_id], device="cuda")
+
+    with torch.no_grad():
+        logits = model(tensor, dataset_ids)
+        probs = F.softmax(logits, dim=-1).squeeze(0).cpu()
+
+    labels = CLASS_LABELS[dataset_name]
+    num_classes = len(labels)
+    # The model pads to C_max; only the first ``num_classes`` entries are real
+    return {labels[i]: float(probs[i]) for i in range(num_classes)}
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+CSS = """
+#col-container { max-width: 900px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.Markdown(
+        "# MoPET: Parameter-Efficient Mixture-of-Experts for Unified Medical Image Classification\n"
+        "Upload a medical image and select a task. The model routes the image through a "
+        "frozen DINOv3 backbone with a sparse mixture-of-experts adapter pool and a "
+        "per-task classification head.\n\n"
+        "[Paper](https://arxiv.org/abs/2607.29462) | "
+        "[GitHub](https://github.com/sdoerrich97/mopet) | "
+        "[Weights](https://huggingface.co/sdoerrich97/mopet_dinov3_unified_blood_breast_derma_path)"
+    )
+
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
+            image_input = gr.Image(type="pil", label="Medical Image", scale=3)
+            dataset_selector = gr.Dropdown(
+                choices=DATASET_NAMES,
+                value=DATASET_NAMES[0],
+                label="Task / Dataset Head",
+                scale=1,
+            )
+
+        run_btn = gr.Button("Classify", variant="primary")
+
+        label_output = gr.Label(num_top_classes=5, label="Classification")
+
+    run_btn.click(
+        fn=classify,
+        inputs=[image_input, dataset_selector],
+        outputs=label_output,
+        api_name="classify",
+    )
+
+    gr.Examples(
+        examples=[
+            ["bloodmnist_sample.png", "BloodMNIST"],
+            ["breastmnist_sample.png", "BreastMNIST"],
+            ["dermamnist_sample.png", "DermaMNIST"],
+            ["pathmnist_sample.png", "PathMNIST"],
+        ],
+        inputs=[image_input, dataset_selector],
+        outputs=label_output,
+        fn=classify,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+demo.launch(mcp_server=True)

mopet/__init__.py

@@ -0,0 +1,19 @@
+"""mopet: parameter-efficient mixture-of-experts for unified medical image classification."""
+
+from __future__ import annotations
+
+from ._factory import create_model, list_pretrained, load_pretrained_weights
+from ._moe import MoEModule, apply_moe_peft
+from .model import MoPET, MultiTaskClassifier
+
+__version__ = "0.1.2"
+__all__ = [
+    "MoPET",
+    "MultiTaskClassifier",
+    "MoEModule",
+    "apply_moe_peft",
+    "create_model",
+    "list_pretrained",
+    "load_pretrained_weights",
+    "__version__",
+]

mopet/_factory.py

@@ -0,0 +1,241 @@
+"""Constructors for MoPET models and pretrained-weight loading.
+
+``create_model`` builds a MoPET model on a timm backbone with the paper's default
+expert configuration; ``load_pretrained_weights`` fetches published checkpoints
+from the HuggingFace Hub. The MedMNIST class-count table is kept here so the
+package can size its heads without importing ``medmnist``.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import cast
+
+import timm
+import torch
+from torch import nn
+
+from .model import MoPET
+
+logger = logging.getLogger(__name__)
+
+#: Friendly backbone name -> timm model id (all ViT-Base/16).
+BACKBONES: dict[str, str] = {
+    "dinov3": "vit_base_patch16_dinov3.lvd1689m",
+    "dino": "vit_base_patch16_224.dino",
+    "clip": "vit_base_patch16_clip_224",
+}
+
+#: Number of classes per MedMNIST+ 2D dataset (used to size the per-dataset heads).
+MEDMNIST_NUM_CLASSES: dict[str, int] = {
+    "BloodMNIST": 8,
+    "BreastMNIST": 2,
+    "ChestMNIST": 14,
+    "DermaMNIST": 7,
+    "OCTMNIST": 4,
+    "OrganAMNIST": 11,
+    "OrganCMNIST": 11,
+    "OrganSMNIST": 11,
+    "PathMNIST": 9,
+    "PneumoniaMNIST": 2,
+    "RetinaMNIST": 5,
+    "TissueMNIST": 8,
+}
+
+#: Default heterogeneous expert pool (the MoPET configuration from the paper).
+DEFAULT_EXPERT_COUNTS: dict[str, int] = {"LoRA": 20, "BOFT": 12}
+DEFAULT_TOP_K: int = 12
+DEFAULT_EXPERT_KWARGS: dict[str, dict[str, object]] = {
+    "LoRA": {"r": 8, "lora_alpha": 8},
+    "BOFT": {"boft_block_size": 8, "boft_n_butterfly_factor": 1},
+    "FourierFT": {"n_frequency": 1000},
+}
+
+
+@dataclass(frozen=True)
+class PublishedModel:
+    """A released MoPET checkpoint: its HuggingFace repo and the exact head layout.
+
+    ``datasets`` is ordered: index i is the dataset id of the i-th classification head,
+    so it must match the order the checkpoint was trained with.
+    """
+
+    repo_id: str
+    backbone: str
+    datasets: tuple[str, ...]
+
+
+#: Released MoPET models (the paper's headline checkpoints). ``create_model(weights=<name>)``
+#: and ``load_pretrained_weights(model, <name>)`` resolve these and pull from the Hub.
+PUBLISHED_MODELS: dict[str, PublishedModel] = {
+    # 4-dataset unified model (Table 2) — also the Breast-booster (Table 3, same pool).
+    "unified": PublishedModel(
+        "sdoerrich97/mopet_dinov3_unified_blood_breast_derma_path",
+        "dinov3",
+        ("BloodMNIST", "BreastMNIST", "DermaMNIST", "PathMNIST"),
+    ),
+    # Auxiliary-booster models (Table 3): a target co-trained with a hand-picked pool.
+    "booster-retina": PublishedModel(
+        "sdoerrich97/mopet_dinov3_booster_retina_breast_blood_retina_path_organa",
+        "dinov3",
+        ("BreastMNIST", "BloodMNIST", "RetinaMNIST", "PathMNIST", "OrganAMNIST"),
+    ),
+    "booster-derma": PublishedModel(
+        "sdoerrich97/mopet_dinov3_booster_derma_derma_blood_oct_organs",
+        "dinov3",
+        ("DermaMNIST", "BloodMNIST", "OCTMNIST", "OrganSMNIST"),
+    ),
+}
+
+
+def list_pretrained() -> dict[str, dict[str, object]]:
+    """List the released MoPET checkpoints and how to load them.
+
+    Returns:
+        A mapping from each variant name (the string passed as ``create_model(weights=...)``)
+        to its ``backbone``, ordered ``datasets`` (index i is the dataset id of head i), and
+        HuggingFace ``repo_id``. Use it to discover the available weights, e.g.::
+
+            import mopet
+            for name, info in mopet.list_pretrained().items():
+                print(name, info["datasets"])
+    """
+    return {
+        name: {
+            "backbone": pub.backbone,
+            "datasets": list(pub.datasets),
+            "repo_id": pub.repo_id,
+        }
+        for name, pub in PUBLISHED_MODELS.items()
+    }
+
+
+def _resolve_backbone(backbone: str) -> str:
+    """Map a friendly backbone name to its timm id (pass-through if already an id)."""
+    return BACKBONES.get(backbone, backbone)
+
+
+def _resolve_num_classes(datasets: list[str] | None, num_classes: list[int] | None) -> list[int]:
+    """Resolve the per-dataset class counts from dataset names or an explicit list."""
+    if num_classes is not None:
+        return num_classes
+    if datasets is not None:
+        try:
+            return [MEDMNIST_NUM_CLASSES[d] for d in datasets]
+        except KeyError as exc:  # pragma: no cover - defensive
+            raise KeyError(f"Unknown MedMNIST dataset: {exc.args[0]!r}") from exc
+    raise ValueError("Provide either `datasets` or `num_classes` to size the heads.")
+
+
+def create_model(
+    backbone: str = "dinov3",
+    datasets: list[str] | None = None,
+    num_classes: list[int] | None = None,
+    pretrained_backbone: bool = True,
+    weights: str | None = None,
+    expert_counts: dict[str, int] | None = None,
+    top_k: int = DEFAULT_TOP_K,
+    expert_kwargs: dict[str, dict[str, object]] | None = None,
+    controller_noise: bool = False,
+    map_location: str = "cpu",
+) -> MoPET:
+    """Build a MoPET model.
+
+    Args:
+        backbone: Friendly name (``"dinov3"``/``"dino"``/``"clip"``) or a timm id.
+        datasets: MedMNIST dataset names defining the multi-task heads (order matters).
+        num_classes: Explicit class counts per head; overrides ``datasets``.
+        pretrained_backbone: Load timm pretrained backbone weights.
+        weights: Either a published-variant name (a key of ``PUBLISHED_MODELS``, e.g.
+            ``"unified"``) to download from the Hub, or a path to a local MoPET
+            checkpoint. A variant name also fixes the backbone and head layout.
+        expert_counts: Experts per family; defaults to the paper's ``{LoRA:20, BOFT:12}``.
+        top_k: Experts activated per token.
+        expert_kwargs: Per-family peft keyword arguments; defaults to the paper's.
+        controller_noise: Whether routers add exploration noise in training.
+        map_location: Device mapping used when loading ``weights``.
+
+    Returns:
+        The constructed :class:`~mopet.model.MoPET`.
+    """
+    published = PUBLISHED_MODELS.get(weights) if weights is not None else None
+    if published is not None:
+        # A released variant fixes the backbone and head layout. Published checkpoints carry
+        # only the trainable parameters (adapters, router, heads); the frozen backbone is
+        # reconstructed from the timm pretrained weights, so keep `pretrained_backbone=True`.
+        backbone = published.backbone
+        datasets = list(published.datasets)
+        num_classes = None
+        pretrained_backbone = True
+
+    timm_id = _resolve_backbone(backbone)
+    heads = _resolve_num_classes(datasets, num_classes)
+    backbone_module = cast(
+        nn.Module, timm.create_model(timm_id, pretrained=pretrained_backbone, num_classes=0)
+    )
+
+    model = MoPET(
+        backbone=backbone_module,
+        num_classes=heads,
+        expert_counts=expert_counts or dict(DEFAULT_EXPERT_COUNTS),
+        top_k=top_k,
+        expert_kwargs=expert_kwargs or {k: dict(v) for k, v in DEFAULT_EXPERT_KWARGS.items()},
+        controller_noise=controller_noise,
+    )
+    if published is not None:
+        _load_into(model, _load_state_dict(_download_variant(published), map_location=map_location))
+    elif weights is not None:
+        _load_into(model, _load_state_dict(weights, map_location=map_location))
+    return model
+
+
+def load_pretrained_weights(model: MoPET, variant: str, map_location: str = "cpu") -> MoPET:
+    """Download and load a published MoPET checkpoint from the HuggingFace Hub.
+
+    Args:
+        model: A MoPET model whose head layout matches ``variant`` (build it with the same
+            ``datasets``/``backbone``, e.g. via ``create_model(weights=variant)``).
+        variant: A key of ``PUBLISHED_MODELS`` (e.g. ``"unified"``, ``"booster-retina"``).
+        map_location: Device mapping for the loaded tensors.
+
+    Returns:
+        ``model`` with the checkpoint loaded in place.
+    """
+    if variant not in PUBLISHED_MODELS:
+        raise KeyError(
+            f"No published weights for {variant!r}. Available: {sorted(PUBLISHED_MODELS)}."
+        )
+    _load_into(model, _load_state_dict(_download_variant(PUBLISHED_MODELS[variant]), map_location))
+    return model
+
+
+def _download_variant(published: PublishedModel) -> str:
+    """Download a published variant's weights from the Hub, preferring safetensors."""
+    from huggingface_hub import hf_hub_download
+
+    try:
+        return hf_hub_download(published.repo_id, filename="model.safetensors")
+    except Exception:  # noqa: BLE001 - fall back to a torch checkpoint
+        return hf_hub_download(published.repo_id, filename="model.pth")
+
+
+def _load_state_dict(path: str, map_location: str = "cpu") -> dict[str, torch.Tensor]:
+    """Load a state dict from a ``.safetensors`` or torch checkpoint file."""
+    if path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+
+        return load_file(path, device=map_location)
+    obj = torch.load(path, map_location=map_location, weights_only=False)
+    state_dict = obj.get("state_dict", obj) if isinstance(obj, dict) else obj
+    return {k.removeprefix("module."): v for k, v in state_dict.items()}
+
+
+def _load_into(model: MoPET, state_dict: dict[str, torch.Tensor]) -> None:
+    """Load trainable parameters into ``model``, tolerating the frozen backbone gap."""
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if unexpected:
+        logger.warning("Unexpected keys when loading MoPET weights: %s", unexpected[:8])
+    logger.info(
+        "Loaded MoPET weights (%d missing, %d unexpected keys).", len(missing), len(unexpected)
+    )

mopet/_moe.py

@@ -0,0 +1,206 @@
+"""Sparse mixture-of-experts over parameter-efficient adapters.
+
+The core of *MoPET*: a learnable top-k router dispatches each token to a small
+subset of a heterogeneous pool of PEFT experts (LoRA / BOFT / FourierFT) that
+wrap a single frozen projection. Each expert returns the full projection output
+(frozen base plus its low-rank delta); because the router weights are a softmax
+over the selected experts, their sum reduces to the frozen projection plus a
+convex combination of the active experts' deltas.
+
+This module is deliberately free of any dataset or configuration framework
+dependency: expert counts, the routing width, and the per-family adapter
+hyperparameters are passed in explicitly.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch.nn.functional as F
+from peft.tuners.boft.layer import Linear as BOFTLinear
+from peft.tuners.fourierft import FourierFTLinear
+from peft.tuners.lora import Linear as LoRALinear
+from torch import nn
+
+logger = logging.getLogger(__name__)
+
+#: Maps an expert-family name to the peft layer class that implements it. The
+#: internal ``peft`` layer classes are used directly (as in the original thesis
+#: code); this is why ``peft`` is pinned exactly.
+_EXPERT_CLASSES: dict[str, type[nn.Module]] = {
+    "LoRA": LoRALinear,
+    "BOFT": BOFTLinear,
+    "FourierFT": FourierFTLinear,
+}
+
+
+class LinearTopKGating(nn.Module):
+    """Linear router producing per-expert routing logits for each token.
+
+    Args:
+        input_dim: Token feature dimension.
+        num_experts: Size of the expert pool to route over.
+        noisy: If ``True``, additive standard-normal noise is applied to the
+            logits during training only, to encourage exploration.
+    """
+
+    def __init__(self, input_dim: int, num_experts: int, noisy: bool = False) -> None:
+        super().__init__()
+        self.num_experts = num_experts
+        self.noisy = noisy
+        # g(x) = W x, no bias.
+        self.gate = nn.Linear(input_dim, num_experts, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Return unnormalized routing logits of shape ``(..., num_experts)``."""
+        logits = self.gate(x)
+        if self.noisy and self.training:
+            logits = logits + torch.randn_like(logits)
+        return logits
+
+
+class MoEModule(nn.Module):
+    """Replaces a single frozen projection with a routed pool of PEFT experts.
+
+    Injected in place of an attention ``qkv`` projection. The wrapped frozen
+    linear is shared as the base layer of every expert, so each expert output is
+    ``base(x) + delta_i(x)`` and the routed, softmax-weighted sum is
+    ``base(x) + sum_i w_i * delta_i(x)`` over the active experts.
+
+    Args:
+        base_linear: The frozen projection to adapt (e.g. attention ``qkv``).
+        expert_counts: Number of experts per family, e.g. ``{"LoRA": 20, "BOFT": 12}``.
+        top_k: Number of experts activated per token.
+        expert_kwargs: Per-family keyword arguments forwarded to the peft layer,
+            e.g. ``{"LoRA": {"r": 8, "lora_alpha": 8}, "BOFT": {...}}``.
+        controller_noise: Whether the router adds exploration noise in training.
+        adapter_name: Adapter name handed to the peft layers.
+    """
+
+    def __init__(
+        self,
+        base_linear: nn.Linear,
+        expert_counts: dict[str, int],
+        top_k: int,
+        expert_kwargs: dict[str, dict[str, object]],
+        controller_noise: bool = False,
+        adapter_name: str = "default",
+    ) -> None:
+        super().__init__()
+        self.top_k = top_k
+        in_dim = base_linear.in_features
+        self.out_features = base_linear.out_features
+
+        self.experts = nn.ModuleList()
+        for family, count in expert_counts.items():
+            if family not in _EXPERT_CLASSES:
+                raise ValueError(f"Unsupported expert family: {family!r}")
+            expert_cls = _EXPERT_CLASSES[family]
+            kwargs = expert_kwargs.get(family, {})
+            for _ in range(count):
+                self.experts.append(
+                    expert_cls(base_layer=base_linear, adapter_name=adapter_name, **kwargs)
+                )
+
+        self.num_experts = len(self.experts)
+        self.controller = LinearTopKGating(
+            input_dim=in_dim, num_experts=self.num_experts, noisy=controller_noise
+        )
+        self._aux_loss: torch.Tensor | float = 0.0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Route tokens through the top-k experts.
+
+        Args:
+            x: Token features of shape ``(B, T, D)``.
+
+        Returns:
+            Adapted projection output of shape ``(B, T, out_features)``.
+        """
+        b, t, _ = x.shape
+        logits = self.controller(x)  # (B, T, E)
+
+        topk_vals, topk_idx = torch.topk(logits, self.top_k, dim=-1)  # (B, T, k)
+        topk_weights = F.softmax(topk_vals, dim=-1)  # (B, T, k)
+
+        output = torch.zeros(b, t, self.out_features, device=x.device, dtype=x.dtype)
+
+        for expert_id, expert in enumerate(self.experts):
+            mask = topk_idx == expert_id  # (B, T, k)
+            if not mask.any():
+                continue
+            b_idx, t_idx, k_idx = mask.nonzero(as_tuple=True)
+            routed_x = x[b_idx, t_idx]  # (N, D)
+            expert_out = expert(routed_x)  # (N, out_features)
+            weights = topk_weights[b_idx, t_idx, k_idx].unsqueeze(-1)  # (N, 1)
+            output[b_idx, t_idx] += expert_out * weights
+
+        self._aux_loss = self.load_balancing_loss(logits, topk_idx)
+        return output
+
+    def load_balancing_loss(self, logits: torch.Tensor, topk_idx: torch.Tensor) -> torch.Tensor:
+        """DeepSeekMoE/Switch-style load-balancing loss.
+
+        Args:
+            logits: Router logits of shape ``(B, T, E)``.
+            topk_idx: Selected expert indices of shape ``(B, T, k)``.
+
+        Returns:
+            Scalar load-balancing loss ``E * sum_i importance_i * load_i``.
+        """
+        b, t, num_experts = logits.shape
+        k = topk_idx.shape[-1]
+
+        probs = torch.softmax(logits, dim=-1)  # (B, T, E)
+        importance = probs.mean(dim=(0, 1))  # (E,)
+
+        one_hot = F.one_hot(topk_idx, num_classes=num_experts)  # (B, T, k, E)
+        load = one_hot.sum(dim=(0, 1, 2)).float() / (b * t * k)  # (E,)
+
+        return num_experts * torch.sum(importance * load)
+
+    def get_aux_loss(self) -> torch.Tensor | float:
+        """Return the load-balancing loss from the most recent forward pass."""
+        return self._aux_loss
+
+
+def apply_moe_peft(
+    model: nn.Module,
+    expert_counts: dict[str, int],
+    top_k: int,
+    expert_kwargs: dict[str, dict[str, object]],
+    controller_noise: bool = False,
+) -> nn.Module:
+    """Replace every attention ``qkv`` projection in ``model`` with a ``MoEModule``.
+
+    Args:
+        model: A timm vision transformer (modified in place).
+        expert_counts: Number of experts per family.
+        top_k: Number of experts activated per token.
+        expert_kwargs: Per-family peft keyword arguments.
+        controller_noise: Whether routers add exploration noise in training.
+
+    Returns:
+        The same ``model``, with its ``qkv`` layers swapped for routed experts.
+    """
+    named_modules = dict(model.named_modules())
+    for name, module in list(model.named_modules()):
+        if not name.endswith("qkv"):
+            continue
+        if "." in name:
+            parent_name, child_name = name.rsplit(".", 1)
+            parent = named_modules[parent_name]
+        else:
+            parent, child_name = model, name
+
+        moe_layer = MoEModule(
+            base_linear=module,
+            expert_counts=expert_counts,
+            top_k=top_k,
+            expert_kwargs=expert_kwargs,
+            controller_noise=controller_noise,
+        )
+        setattr(parent, child_name, moe_layer)
+
+    return model

mopet/model.py

@@ -0,0 +1,119 @@
+"""The MoPET model: a frozen foundation backbone adapted by a routed PEFT expert pool.
+
+``MoPET`` freezes a timm vision-transformer backbone, replaces its attention
+``qkv`` projections with sparse mixture-of-experts adapters (see :mod:`mopet._moe`),
+and attaches one classification head per dataset so a single network serves many
+heterogeneous medical-image classification tasks at once.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import cast
+
+import torch
+from torch import nn
+
+from ._moe import MoEModule, apply_moe_peft
+
+logger = logging.getLogger(__name__)
+
+
+class MultiTaskClassifier(nn.Module):
+    """One linear classification head per dataset with padded, per-sample routing.
+
+    Args:
+        num_classes: Number of classes for each dataset, in a fixed order; the
+            index into this list is the dataset identifier used at forward time.
+        input_dim: Dimension of the shared backbone feature.
+    """
+
+    def __init__(self, num_classes: list[int], input_dim: int) -> None:
+        super().__init__()
+        self.classifiers = nn.ModuleList(
+            nn.Linear(in_features=input_dim, out_features=n) for n in num_classes
+        )
+
+    def forward(self, x: torch.Tensor, dataset_ids: torch.Tensor) -> torch.Tensor:
+        """Route each sample to its dataset head.
+
+        Args:
+            x: Shared features of shape ``(B, D)``.
+            dataset_ids: Dataset index per sample, shape ``(B,)``.
+
+        Returns:
+            Logits of shape ``(B, C_max)`` where ``C_max`` is the largest class
+            count present in the batch; unused entries are padded with ``-1e9``.
+        """
+        b = x.size(0)
+        present: list[int] = torch.unique(dataset_ids).tolist()
+        heads = [cast(nn.Linear, self.classifiers[task]) for task in present]
+        max_classes = max(head.out_features for head in heads)
+
+        output = torch.full((b, max_classes), fill_value=-1e9, device=x.device, dtype=x.dtype)
+        for task, head in zip(present, heads, strict=True):
+            idx = (dataset_ids == task).nonzero(as_tuple=True)[0]
+            logits = head(x[idx])
+            output[idx, : logits.size(1)] = logits
+        return output
+
+
+class MoPET(nn.Module):
+    """Frozen backbone + routed PEFT experts + per-dataset heads.
+
+    Args:
+        backbone: A timm vision transformer providing ``(B, D)`` pooled features
+            once its own head is removed. Frozen in place.
+        num_classes: Class count per dataset (defines the multi-task heads).
+        expert_counts: Number of experts per family, e.g. ``{"LoRA": 20, "BOFT": 12}``.
+        top_k: Number of experts activated per token.
+        expert_kwargs: Per-family peft keyword arguments.
+        controller_noise: Whether routers add exploration noise in training.
+    """
+
+    def __init__(
+        self,
+        backbone: nn.Module,
+        num_classes: list[int],
+        expert_counts: dict[str, int],
+        top_k: int,
+        expert_kwargs: dict[str, dict[str, object]],
+        controller_noise: bool = False,
+    ) -> None:
+        super().__init__()
+        self.pretrained_cfg = getattr(backbone, "pretrained_cfg", None)
+
+        backbone.head = nn.Identity()
+        for param in backbone.parameters():
+            param.requires_grad = False
+
+        self.backbone = apply_moe_peft(
+            backbone,
+            expert_counts=expert_counts,
+            top_k=top_k,
+            expert_kwargs=expert_kwargs,
+            controller_noise=controller_noise,
+        )
+        input_dim = int(cast(int, backbone.num_features))
+        self.head = MultiTaskClassifier(num_classes=num_classes, input_dim=input_dim)
+
+    def forward(self, x: torch.Tensor, dataset_ids: torch.Tensor) -> torch.Tensor:
+        """Classify a batch of images tagged with their dataset ids.
+
+        Args:
+            x: Input images of shape ``(B, C, H, W)``.
+            dataset_ids: Dataset index per sample, shape ``(B,)``.
+
+        Returns:
+            Padded per-dataset logits of shape ``(B, C_max)``.
+        """
+        features = self.backbone(x)
+        return self.head(features, dataset_ids)
+
+    def get_aux_loss(self) -> torch.Tensor | float:
+        """Sum the load-balancing loss over all routed layers from the last forward."""
+        aux_loss: torch.Tensor | float = 0.0
+        for module in self.backbone.modules():
+            if isinstance(module, MoEModule):
+                aux_loss = aux_loss + module.get_aux_loss()
+        return aux_loss

requirements.txt

@@ -0,0 +1,5 @@
+torchvision
+timm>=1.0.22,<2
+peft>=0.18,<0.19
+safetensors>=0.5
+pillow>=10.0