rigfl 0.1.0__py3-none-any.whl

rigfl/__init__.py

@@ -0,0 +1,8 @@
+"""RigFL — a modular framework for rigorous federated learning experimentation."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("rigfl")
+except PackageNotFoundError:
+    __version__ = "0+unknown"

rigfl/algorithms/__init__.py

@@ -0,0 +1 @@
+"""Federated algorithms. Each file is one self-contained algorithm."""

rigfl/algorithms/fedavg.py

@@ -0,0 +1,164 @@
+"""FedAvg -- homogeneous full-model federated averaging.
+
+Every client starts a round from the same global ``ClientModel``, trains it with
+local SGD, and uploads a detached state snapshot plus its local sample count.
+The server averages floating-point state by sample count. Non-floating buffers
+cannot be meaningfully averaged, so they are copied from the largest upload
+(ties follow client/upload order).
+"""
+
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass
+from typing import Mapping
+
+import torch
+import torch.nn.functional as F
+from pydantic import Field
+
+from rigfl.core.config import AlgorithmConfig
+from rigfl.core.interfaces import Algorithm
+from rigfl.prediction import Predictions
+
+
+class FedAvgConfig(AlgorithmConfig):
+    local_epochs: int = Field(1, ge=1)
+    lr: float = Field(0.01, gt=0)
+
+
+@dataclass(frozen=True)
+class ModelUpload:
+    """A complete locally trained model state and its aggregation weight."""
+
+    state: dict[str, torch.Tensor]
+    num_samples: int
+
+
+def clone_state_dict(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """A detached snapshot: later client training cannot mutate the upload."""
+    return {name: value.detach().clone() for name, value in state.items()}
+
+
+def validate_state_structure(local, global_state: Mapping[str, torch.Tensor], *,
+                             algorithm: str, client_id=None) -> None:
+    """Reject parameter/buffer layouts that full-model FL cannot synchronize."""
+    local_state = local.state_dict() if hasattr(local, "state_dict") else local
+    local_keys, global_keys = set(local_state), set(global_state)
+    where = f" client {client_id}" if client_id is not None else ""
+    prefix = f"{algorithm} requires homogeneous client models;{where}"
+    if local_keys != global_keys:
+        missing = sorted(global_keys - local_keys)
+        extra = sorted(local_keys - global_keys)
+        detail = []
+        if missing:
+            detail.append(f"missing keys {missing}")
+        if extra:
+            detail.append(f"extra keys {extra}")
+        raise ValueError(f"{prefix} has an incompatible state structure ({'; '.join(detail)}).")
+    for name, reference in global_state.items():
+        value = local_state[name]
+        if value.shape != reference.shape:
+            raise ValueError(
+                f"{prefix} state {name!r} has shape {tuple(value.shape)}, expected "
+                f"{tuple(reference.shape)}."
+            )
+        if value.dtype != reference.dtype:
+            raise ValueError(
+                f"{prefix} state {name!r} has dtype {value.dtype}, expected "
+                f"{reference.dtype}."
+            )
+
+
+def weighted_average_states(uploads: list[ModelUpload], *, device,
+                            algorithm: str = "FedAvg") -> dict[str, torch.Tensor]:
+    """Sample-count-weighted floating state plus deterministic integer buffers."""
+    if not uploads:
+        raise ValueError("FedAvg aggregation requires at least one client upload.")
+    if any(upload.num_samples < 0 for upload in uploads):
+        raise ValueError("FedAvg upload sample counts must be non-negative.")
+    total = sum(upload.num_samples for upload in uploads)
+    if total <= 0:
+        raise ValueError("FedAvg aggregation requires at least one training sample.")
+
+    reference = uploads[0].state
+    for upload in uploads[1:]:
+        validate_state_structure(upload.state, reference, algorithm=f"{algorithm} upload")
+
+    # There is no arithmetic mean for integer/categorical module state. Selecting
+    # the largest-weight contributor preserves dtype and a real client value.
+    source = max(enumerate(uploads), key=lambda item: (item[1].num_samples, -item[0]))[1]
+    averaged: dict[str, torch.Tensor] = {}
+    for name, first in reference.items():
+        if torch.is_floating_point(first) or torch.is_complex(first):
+            if torch.is_complex(first):
+                work_dtype = torch.complex128 if first.dtype == torch.complex128 else torch.complex64
+            else:
+                work_dtype = torch.float64 if first.dtype == torch.float64 else torch.float32
+            value = torch.zeros(first.shape, dtype=work_dtype, device=device)
+            for upload in uploads:
+                value.add_(upload.state[name].to(device=device, dtype=work_dtype),
+                           alpha=upload.num_samples / total)
+            averaged[name] = value.to(dtype=first.dtype)
+        else:
+            averaged[name] = source.state[name].to(device=device).clone()
+    return averaged
+
+
+class FedAvg(Algorithm):
+    """Traditional FedAvg over the complete homogeneous client model."""
+
+    algorithm_name = "FedAvg"
+
+    def __init__(self, config: FedAvgConfig, model_template):
+        super().__init__(config)
+        if model_template is None:
+            raise ValueError("FedAvg requires an initial homogeneous client model template.")
+        self.model_template = copy.deepcopy(model_template).cpu()
+
+    @classmethod
+    def from_config(cls, config, *, model_template=None, **resources):
+        return cls(config, model_template)
+
+    def init_globals(self):
+        return copy.deepcopy(self.model_template)
+
+    def _round_reference(self, model) -> dict[str, torch.Tensor] | None:
+        return None
+
+    def _training_loss(self, model, logits, labels, reference) -> torch.Tensor:
+        return F.cross_entropy(logits, labels)
+
+    def local_train(self, client, global_model) -> ModelUpload:
+        model, loader = client.model, client.train_loader
+        global_state = global_model.state_dict()
+        validate_state_structure(model, global_state, algorithm=self.algorithm_name,
+                                 client_id=client.client_id)
+        model.load_state_dict(global_state)
+        model.to(self.device)
+        model.train()
+        reference = self._round_reference(model)
+        optimizer = torch.optim.SGD(model.parameters(), lr=self.config.lr)
+        for _ in range(self.config.local_epochs):
+            for x, y in loader:
+                x, y = x.to(self.device), y.to(self.device)
+                loss = self._training_loss(model, model(x), y, reference)
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+        return ModelUpload(clone_state_dict(model.state_dict()), len(loader.dataset))
+
+    def aggregate(self, uploads: list[ModelUpload], global_model):
+        averaged = weighted_average_states(
+            uploads, device=self.device, algorithm=self.algorithm_name
+        )
+        validate_state_structure(averaged, global_model.state_dict(),
+                                 algorithm=self.algorithm_name)
+        global_model.to(self.device)
+        global_model.load_state_dict(averaged)
+        return global_model
+
+    @torch.no_grad()
+    def predict(self, client, x, global_model) -> Predictions:
+        global_model.to(x.device)
+        return Predictions.from_logits(global_model(x))

rigfl/algorithms/feddes.py

@@ -0,0 +1,380 @@
+"""FedDES -- Federated Diverse Ensemble Selection (Mueller & Street).
+
+FedDES uses RigFL's peer-to-peer one-shot lifecycle:
+    prepare                  each client trains its local classifier pool.
+    one_shot_communication   the local pools are shared with every client once.
+    local_computation        each client independently builds its graph and
+                             trains its complete GNN meta-learner.
+
+There is no iterative local-training/server-aggregation loop. Each GNN retains
+the epoch selected by its own validation split, and RigFL evaluates those final
+per-client models once.
+
+Per-client GraphRoute state is stored in ``ctx.client_state``. GraphRoute remains
+an optional dependency for users running other algorithms.
+"""
+
+from __future__ import annotations
+
+import copy
+from pathlib import Path
+from typing import Callable, Literal
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from pydantic import Field
+
+from rigfl.core.config import AlgorithmConfig
+from rigfl.core.interfaces import Algorithm, LocalSelection, OneShotContext
+from rigfl.prediction import Predictions
+from rigfl.data.builder import _collate    # multi-input-safe collate (works for single-input too)
+
+FEDDES_PREPROCESSING_KEY = "rigfl-feddes-multitensor-collation-v1"
+
+
+class FedDESConfig(AlgorithmConfig):
+    # FedDES trains base classifiers + a GNN meta-learner; it does not use the
+    # local_epochs/lr settings used by algorithms with a client training loop.
+    gnn_arch: Literal["gat", "graph_gps", "mlp"] = "gat"
+    base_lr: float = Field(5e-4, gt=0)
+    base_epochs: int = Field(100, ge=1)
+    graph_k: int = Field(5, ge=1)
+    hidden_dim: int = Field(128, ge=1)
+    gnn_epochs: int = Field(500, ge=1)
+    gnn_patience: int = Field(50, ge=1)
+    calibrate: bool = True
+    # OOF stacking produces training meta-labels from models that did not see the
+    # corresponding rows. In-sample mode is cheaper but measures training fit.
+    base_split_mode: Literal["oof_stacking", "in_sample"] = "oof_stacking"
+    base_oof_folds: int = Field(3, ge=2)
+    cache_dir: str = "pool_cache"     # reuse trained base pools across graph/GNN sweeps ("" disables)
+
+
+class FedDES(Algorithm):
+    def __init__(self, config: FedDESConfig,
+                 base_factories: list[torch.nn.Module | Callable[[], torch.nn.Module]],
+                 num_classes: int, *, data_id: str | None = None,
+                 model_ids: list[str] | None = None, seed: int = 0,
+                 validation_fraction: float = 0.2):
+        super().__init__(config)
+        sources = list(base_factories)
+        # Configuration resolves to actual model templates. A few low-level test
+        # integrations still supply constructors, so materialize those once and
+        # use the resulting templates for both identity and isolated training.
+        self.base_models = tuple(
+            source if isinstance(source, torch.nn.Module) else source()
+            for source in sources
+        )
+        self.base_factories = [
+            lambda template=model: copy.deepcopy(template)
+            for model in self.base_models
+        ]
+        self.num_classes = num_classes
+        self.gnn_arch = config.gnn_arch
+        self.base_lr, self.base_epochs = config.base_lr, config.base_epochs
+        self.graph_k, self.hidden_dim = config.graph_k, config.hidden_dim
+        self.calibrate = config.calibrate
+        self.base_split_mode = config.base_split_mode
+        self.base_oof_folds = config.base_oof_folds
+        self.cache_dir = config.cache_dir or None
+        self.data_id, self.seed = data_id, seed
+        self.validation_fraction = validation_fraction
+        self.model_ids = model_ids
+        if self.cache_dir and not self.data_id:
+            raise ValueError("FedDES pool reuse requires a stable data_id.")
+        if self.cache_dir and not self.model_ids:
+            raise ValueError("FedDES pool reuse requires stable model_ids.")
+        if self.model_ids and len(self.model_ids) != len(self.base_models):
+            raise ValueError("model_ids must name every base model in order.")
+        if self.model_ids is None:
+            self.model_ids = [f"model_{i}" for i in range(len(base_factories))]
+        self.gnn_epochs, self.gnn_patience = config.gnn_epochs, config.gnn_patience
+
+    @classmethod
+    def from_config(cls, config, *, experiment, base_pool=None,
+                    model_input_spec=None, **resources):
+        from rigfl.core.adapters import LearnedProjection
+        from rigfl.experiment.config import fingerprint
+        from rigfl.models.registry import (instantiate_models,
+                                           resolve_model_architectures)
+
+        if experiment.scheme in {"generated", "natural"}:
+            data_id = f"{experiment.dataset}-{experiment.partition}"
+        else:
+            partition_settings = {
+                "dataset": experiment.dataset,
+                "num_clients": experiment.num_clients,
+                "alpha": experiment.alpha,
+                "seed": experiment.seed,
+                "train_per_client": experiment.train_per_client,
+                "test_per_client": experiment.test_per_client,
+                "val_frac": experiment.val_frac,
+            }
+            data_id = (
+                f"{experiment.dataset}-{fingerprint(partition_settings)}")
+
+        input_kind = model_input_spec["input_kind"] if model_input_spec else (
+            "temporal" if experiment.scheme == "natural" else "image")
+        model_ids = resolve_model_architectures(
+            architecture_family=experiment.model_architecture_family,
+            architectures=experiment.model_architectures,
+            input_kind=input_kind,
+        )
+        if base_pool is None:
+            if model_input_spec is None:
+                if experiment.scheme == "natural":
+                    raise ValueError(
+                        "FedDES temporal models require model_input_spec with "
+                        "n_ts, n_static, and seq_len."
+                    )
+                model_input_spec = {
+                    "input_kind": "image", "shape": (3, 32, 32)
+                }
+            base_pool = instantiate_models(
+                model_ids,
+                num_classes=experiment.num_classes,
+                input_spec=model_input_spec,
+                shared_dim=experiment.shared_dim,
+                adapter=lambda native, shared: LearnedProjection(native, shared),
+            )
+        return cls(
+            config,
+            base_pool,
+            experiment.num_classes,
+            data_id=data_id,
+            model_ids=model_ids,
+            seed=experiment.seed,
+            validation_fraction=experiment.val_frac,
+        )
+
+    def prepare(self, model, train_loader, ctx: OneShotContext):
+        """Train one client's local base pool and return its outgoing payload."""
+        del model
+        if ctx.validation_loader is None:
+            raise ValueError("FedDES requires each client's official validation split.")
+        st = ctx.client_state
+        st["train_dataset"] = train_loader.dataset
+        st["validation_dataset"] = ctx.validation_loader.dataset
+        st["local_pool"] = self._train_or_load_pool(
+            st["train_dataset"], st["validation_dataset"],
+            ctx.device, ctx.client_id)
+        return st["local_pool"]
+
+    def one_shot_communication(self, outgoing: list):
+        """Share the ordered union of all local pools with every client once."""
+        shared_pool = self._union_pools(outgoing)
+        return [shared_pool for _ in outgoing]
+
+    # ── decision space: pool probabilities [N,M*C] + per-classifier hard preds [N,M] ──
+    def _splice_oof(self, tr_logits, pool, st, client_id):
+        """Replace this client's own columns with their out-of-fold logits.
+
+        The global pool is every client's classifiers concatenated, and only the
+        slice this client contributed was fit on tr_ds -- the others already
+        predict it out-of-sample, since the rows are not their data. Explicit
+        client/model IDs locate the local columns in the ordered global pool.
+        """
+        local_pool = st.get("local_pool")
+        oof = None if local_pool is None else local_pool.load_oof()
+        if oof is None:
+            return tr_logits
+        prefix = f"client_{client_id}/"
+        cols = [j for j, model_id in enumerate(pool.model_ids)
+                if model_id.startswith(prefix)]
+        if len(cols) != oof.shape[1] or oof.shape[0] != tr_logits.shape[0]:
+            print(f"[FedDES][warn] cannot place OOF logits (matched {len(cols)} of "
+                  f"{oof.shape[1]} own classifiers, {oof.shape[0]} vs "
+                  f"{tr_logits.shape[0]} rows); using in-sample meta-labels.")
+            return tr_logits
+        tr_logits[:, cols, :] = oof.to(tr_logits.device, tr_logits.dtype)
+        return tr_logits
+
+    # ── base-pool artifact reuse (train once; reuse across graph/GNN sweeps) ──
+    def _pool_fp(self) -> str:
+        """Fingerprint the ordered local pool and its complete training policy."""
+        from graphroute.pool_cache import fingerprint_model, fingerprint_pool
+        from rigfl.experiment.env import _package
+        template_fingerprints = [
+            fingerprint_model(model) for model in self.base_models
+        ]
+        return fingerprint_pool(
+            model_ids=self.model_ids,
+            model_fingerprints=template_fingerprints,
+            base_config={
+                "task": "classification", "num_classes": self.num_classes,
+                "split_mode": self.base_split_mode,
+                "oof_folds": self.base_oof_folds,
+                "lr": self.base_lr, "epochs": self.base_epochs,
+                "batch_size": 64, "patience": 20, "optimizer": "Adam",
+                "weight_decay": 5e-4, "weighted_by_class": True,
+                "es_metric": "val_loss", "inner_val_ratio": 0.2,
+                "client_validation_fraction": self.validation_fraction,
+                "seed_policy": "experiment_seed_plus_client_id",
+                "preprocessing": FEDDES_PREPROCESSING_KEY,
+            },
+            seed=self.seed,
+            code_identity={"graphroute": _package("graphroute"),
+                           "rigfl": _package("rigfl")})
+
+    def _train(self, tr_ds, va_ds, device, client_id):
+        """Return the trained models and optional out-of-fold logits."""
+        from graphroute.run import seed_everything
+        seed_everything(self.seed + int(client_id))
+        if self.base_split_mode == "oof_stacking":
+            from graphroute.pool import train_pool_oof
+            models, oof_logits, _ = train_pool_oof(
+                self.base_factories, tr_ds, va_ds, device,
+                n_folds=self.base_oof_folds, num_classes=self.num_classes,
+                lr=self.base_lr, max_epochs=self.base_epochs,
+                seed=self.seed + int(client_id), collate_fn=_collate)
+            return models, oof_logits          # [N_tr, M_local, C], row i unseen by its predictor
+        from graphroute.pool import train_pool
+        return train_pool(self.base_factories, tr_ds, va_ds, device,
+                          num_classes=self.num_classes, lr=self.base_lr, max_epochs=self.base_epochs,
+                          collate_fn=_collate), None   # keep multi-input (ts, static) as a MultiTensor
+
+    def _train_or_load_pool(self, tr_ds, va_ds, device, client_id):
+        """Load or train this client's pool for reuse across graph/GNN sweeps."""
+        def train():
+            return self._train(tr_ds, va_ds, device, client_id)
+
+        if not self.cache_dir:
+            from graphroute.pool_cache import in_memory_pool
+            models, oof = train()
+            artifact = in_memory_pool(
+                models, model_ids=self.model_ids,
+                fingerprint_value=self._pool_fp())
+            artifact.oof_logits = oof
+            return artifact
+        from pathlib import Path
+
+        from graphroute.pool_cache import cached_pool
+        fp = self._pool_fp()
+        print(f"[FedDES] base pool {fp} (client {client_id})")
+        directory = (Path(self.cache_dir) / self.data_id / f"pool_{fp}"
+                     / "clients" / f"client_{client_id}")
+        return cached_pool(
+            directory, self.base_factories, train,
+            fingerprint_value=fp, model_ids=self.model_ids,
+            require_oof=self.base_split_mode == "oof_stacking",
+            data_id=self.data_id)
+
+    def _graphroute_config(self, client_id, device):
+        from graphroute.config import GraphRouteConfig
+        return GraphRouteConfig(
+            task="classification", loss_target="meta_labels",
+            dataset=self.data_id or "federated-client",
+            num_classes=self.num_classes, seed=self.seed + int(client_id),
+            device=device.type,
+            base={"split_mode": ("oof_stacking" if self.base_split_mode == "oof_stacking"
+                                  else "split_train")},
+            graph={"k": self.graph_k, "pool_calibrate": self.calibrate},
+            gnn={"arch": self.gnn_arch, "hidden_dim": self.hidden_dim,
+                 "epochs": self.gnn_epochs, "patience": self.gnn_patience,
+                 "es_metric": "val_acc",
+                 "ens_combination_mode": "hard_weighted_voting",
+                 "voting_weight_space": "sig"})
+
+    # ── post-communication local computation: train one complete local GNN ──
+    def local_computation(self, model, pool, loader, ctx: OneShotContext):
+        del model, loader
+        device, st = ctx.device, ctx.client_state
+        if pool is None:
+            raise RuntimeError("FedDES communication did not provide a classifier pool.")
+        train_dataset = st["train_dataset"]
+        validation_dataset = st["validation_dataset"]
+        tr_loader = DataLoader(train_dataset, 256, shuffle=False, collate_fn=_collate)
+        output_dir = (pool.directory / "outputs" / f"client_{ctx.client_id}"
+                      if pool.directory is not None else None)
+        client_pool = pool.for_data(output_directory=output_dir)
+        train_logits = client_pool.cached_outputs(
+            self._output_name("train"),
+            tr_loader, device, task="classification",
+            transform=lambda value: self._splice_oof(
+                value, pool, st, ctx.client_id))
+        client_pool = client_pool.for_data(
+            output_directory=output_dir, training_outputs=train_logits)
+
+        from graphroute.run import fit_graphroute
+        st["graphroute_model"] = fit_graphroute(
+            self._graphroute_config(ctx.client_id, device), train_dataset,
+            validation_set=validation_dataset, pool=client_pool,
+            collate_fn=_collate)
+        training = st["graphroute_model"].history
+        return LocalSelection(
+            selected_step=int(training["best_epoch"]),
+            metric="accuracy",
+            validation_value=float(training["best_metric"]),
+        )
+
+    @staticmethod
+    def _output_name(split: str) -> str:
+        return f"{split}_logits"
+
+    def _union_pools(self, uploads: list):
+        """Combine the prepared client pools in deterministic client/model order."""
+        from graphroute.pool_cache import PoolArtifact, fingerprint
+        model_ids, factories, paths, models = [], [], [], []
+        can_hold_models = all(artifact.models is not None for artifact in uploads)
+        for cid, artifact in enumerate(uploads):
+            model_ids.extend(f"client_{cid}/{name}" for name in artifact.model_ids)
+            factories.extend(artifact.model_factories)
+            paths.extend(artifact.model_paths)
+            if can_hold_models:
+                models.extend(artifact.models)
+        root = (None if not self.cache_dir else
+                Path(self.cache_dir) / self.data_id / f"pool_{self._pool_fp()}")
+        shared_fingerprint = fingerprint({"ordered_members": model_ids,
+                                          "local_pool": self._pool_fp()})
+        if root is not None:
+            _write_manifest(root, shared_fingerprint, model_ids, self.data_id)
+        return PoolArtifact(
+            fingerprint=shared_fingerprint,
+            model_ids=tuple(model_ids), model_factories=tuple(factories),
+            model_paths=tuple(paths), directory=root,
+            output_directory=None, models=models if can_hold_models else None)
+
+    # prediction: connect the query batch into the train graph, run the GNN, ensemble-select
+    def predict(self, client, x, shared) -> Predictions:
+        st = client.state
+        if "graphroute_model" not in st:
+            raise RuntimeError(
+                "FedDES cannot predict before its local GNN has been trained.")
+
+        predicted = st["graphroute_model"].predict(
+            _BatchDataset(x), split="batch", cache_outputs=False)
+        return Predictions.from_probabilities(
+            predicted["probabilities"], labels=predicted["predictions"])
+
+# ── small helpers ────────────────────────────────────────────────────────────
+class _BatchDataset(Dataset):
+    """Wrap a query batch as a (labelless) Dataset so it goes through the pool.
+    Handles single-input ``x`` (a tensor) and multi-input ``x`` (a MultiTensor /
+    tuple of tensors, e.g. eICU's ``(ts, static)``) -- indexing samples, not fields."""
+    def __init__(self, x):
+        self.x = x
+        self.multi = isinstance(x, tuple)          # MultiTensor is a tuple subclass
+    def __len__(self):
+        return len(self.x[0]) if self.multi else len(self.x)
+    def __getitem__(self, i):
+        return (tuple(f[i] for f in self.x) if self.multi else self.x[i]), 0
+
+
+def _write_manifest(directory: Path, pool_fingerprint: str,
+                    model_ids: list[str], data_id: str) -> None:
+    """Record the deterministic order of the federated pool."""
+    import json
+    import os
+
+    directory.mkdir(parents=True, exist_ok=True)
+    path = directory / "manifest.json"
+    value = {"schema_version": 2, "data_id": data_id,
+             "fingerprint": pool_fingerprint, "model_ids": model_ids}
+    if path.exists():
+        if json.loads(path.read_text()) != value:
+            raise RuntimeError(f"FedDES pool manifest does not match {path}")
+        return
+    tmp = path.with_name(path.name + f".{os.getpid()}.tmp")
+    tmp.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n")
+    os.replace(tmp, path)

rigfl/algorithms/fedgh.py

@@ -0,0 +1,106 @@
+"""FedGH -- Federated Global Header (Yi et al., ACM MM 2023).
+
+Clients keep their feature extractors but share a global classifier head. Each
+round, the server broadcasts the head; each client installs it,
+fine-tunes locally, and uploads per-class prototypes (mean representations);
+the server then *trains* the head on those (prototype -> label) pairs and keeps
+it for the next round. Predictions use the head.
+
+Server-side header training follows Algorithm 1 / Eq. 4: the optimizer updates
+the header from uploaded (representation, label) pairs. See DEVIATIONS.md.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pydantic import Field
+
+from rigfl.core.config import AlgorithmConfig
+from rigfl.core.interfaces import Algorithm
+from rigfl.prediction import Predictions
+
+Prototypes = dict[int, torch.Tensor]
+
+
+class FedGHConfig(AlgorithmConfig):
+    local_epochs: int = Field(1, ge=1)
+    lr: float = Field(0.01, gt=0)
+    server_epochs: int = Field(1, ge=1)
+    server_lr: float = Field(0.01, gt=0)
+
+
+class FedGH(Algorithm):
+    def __init__(self, config: FedGHConfig, shared_dim: int, num_classes: int):
+        super().__init__(config)
+        self.shared_dim = shared_dim
+        self.num_classes = num_classes
+
+    @classmethod
+    def from_config(cls, config, *, experiment, **resources):
+        return cls(config, experiment.shared_dim, experiment.num_classes)
+
+    def init_globals(self) -> nn.Linear:
+        return nn.Linear(self.shared_dim, self.num_classes)   # the shared, server-trained head
+
+    def local_train(self, client, global_head: nn.Linear) -> Prototypes:
+        model, loader = client.model, client.train_loader
+        device = self.device
+        model.head.load_state_dict(global_head.state_dict())   # install the global head
+        model.to(device)
+        model.train()
+        optimizer = torch.optim.SGD(model.parameters(), lr=self.config.lr)
+
+        for _ in range(self.config.local_epochs):
+            for x, y in loader:
+                x, y = x.to(device), y.to(device)
+                loss = F.cross_entropy(model(x), y)
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        return self._local_prototypes(model, loader, device)
+
+    @torch.no_grad()
+    def _local_prototypes(self, model, loader, device) -> Prototypes:
+        model.eval()
+        total: dict[int, torch.Tensor] = {}
+        count: dict[int, int] = defaultdict(int)
+        for x, y in loader:
+            x, y = x.to(device), y.to(device)
+            for r, label in zip(model.rep(x), y):
+                c = label.item()
+                total[c] = r.clone() if c not in total else total[c] + r
+                count[c] += 1
+        return {c: total[c] / count[c] for c in total}
+
+    def aggregate(self, uploads: list[Prototypes],
+                  global_head: nn.Linear) -> nn.Linear:
+        device = self.device
+        pairs = [(proto, c) for protos in uploads for c, proto in protos.items()]
+        global_head = global_head.to(device)
+        global_head.train()
+        optimizer = torch.optim.SGD(
+            global_head.parameters(), lr=self.config.server_lr)
+
+        for _ in range(self.config.server_epochs):
+            for proto, c in pairs:
+                logit = global_head(proto.unsqueeze(0).to(device))
+                loss = F.cross_entropy(logit, torch.tensor([c], device=device))
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        return global_head
+
+    @torch.no_grad()
+    def predict(self, client, x, global_head: nn.Linear) -> Predictions:
+        model = client.model
+        # Local backbone + the server-trained head: the model that actually
+        # predicts, and so the one whose logits become the distribution.
+        model.head.load_state_dict(global_head.state_dict())
+        return Predictions.from_logits(model(x))