Training API¶

Auto-generated from docstrings.

`ecgen.training.runner`¶

YAML-driven training runner shared by all families.

Config sections::

experiment: {name, seed, runs_root}
model:      {name | target, params}          # registry name, e.g. "vae"
data:       {dataset: {name | target, params}, batch_size, num_workers,
             pin_memory, split_params}
checkpoint: {monitor, mode, save_top_k, save_last, filename}
callbacks:  [{name | target, params}, ...]
trainer:    {...}                            # pl.Trainer kwargs
wandb:      {enabled, project, entity, run_name, tags, notes, log_model}
resume:     path/to/checkpoint.ckpt | null

Outputs go to <runs_root>/<name>/seed_<seed>/ (checkpoints/, samples/, metrics/, tb/, config_resolved.yaml, metadata.json). String values may use ${run_dir}, ${name} and ${seed}; a callback that takes a save_dir and gets none writes to ${run_dir}/samples.

parse_value ¶

parse_value(raw: str) -> Any

Parse an override value as YAML, reading exponent notation as a float.

Source code in src/ecgen/training/runner.py

def parse_value(raw: str) -> Any:
    """Parse an override value as YAML, reading exponent notation as a float."""
    value = yaml.safe_load(raw)
    if isinstance(value, str) and _FLOAT_RE.match(value.strip()):
        return float(value)
    return value

apply_overrides ¶

apply_overrides(cfg: Dict[str, Any], overrides: Sequence[str]) -> Dict[str, Any]

Apply dotted.key=value overrides; values are parsed as YAML.

model.params.lr=1e-3, trainer.devices=[1], callbacks.0.params.num_samples=2.

Source code in src/ecgen/training/runner.py

def apply_overrides(cfg: Dict[str, Any], overrides: Sequence[str]) -> Dict[str, Any]:
    """
    Apply ``dotted.key=value`` overrides; values are parsed as YAML.

    ``model.params.lr=1e-3``, ``trainer.devices=[1]``, ``callbacks.0.params.num_samples=2``.
    """
    cfg = copy.deepcopy(cfg)
    for item in overrides:
        if "=" not in item:
            raise ValueError(f"Override must look like key=value, got {item!r}")
        key, raw = item.split("=", 1)
        value = parse_value(raw)
        node: Any = cfg
        parts = key.split(".")
        for part in parts[:-1]:
            if isinstance(node, list):
                node = node[int(part)]
            else:
                node = node.setdefault(part, {})
                if node is None:
                    raise ValueError(f"Cannot set {key!r}: {part!r} is null")
        last = parts[-1]
        if isinstance(node, list):
            node[int(last)] = value
        else:
            node[last] = value
    return cfg

resolve_config ¶

resolve_config(cfg: Mapping[str, Any]) -> tuple[Dict[str, Any], RunPaths]

Validate sections, fill defaults, and interpolate ${...} variables.

Source code in src/ecgen/training/runner.py

def resolve_config(cfg: Mapping[str, Any]) -> tuple[Dict[str, Any], RunPaths]:
    """Validate sections, fill defaults, and interpolate ``${...}`` variables."""
    unknown = set(cfg) - KNOWN_SECTIONS
    if unknown:
        raise ValueError(
            f"Unknown config sections: {sorted(unknown)}; expected {sorted(KNOWN_SECTIONS)}"
        )
    for section in ("model", "data"):
        if section not in cfg:
            raise ValueError(f"Config needs a '{section}' section")

    cfg = copy.deepcopy(dict(cfg))
    exp = cfg.setdefault("experiment", {})
    exp.setdefault("name", "ecg_experiment")
    exp["seed"] = int(exp.get("seed", 42))
    exp.setdefault("runs_root", "runs")

    run_dir = Path(exp["runs_root"]) / exp["name"] / f"seed_{exp['seed']}"
    cfg = _interpolate(
        cfg, {"run_dir": str(run_dir), "name": exp["name"], "seed": str(exp["seed"])}
    )
    cfg["checkpoint"] = {**DEFAULT_CHECKPOINT, **(cfg.get("checkpoint") or {})}
    cfg.setdefault("callbacks", [])
    cfg["callbacks"] = cfg["callbacks"] or []
    cfg.setdefault("trainer", {})
    cfg.setdefault("wandb", {"enabled": False})
    cfg.setdefault("resume", None)

    paths = RunPaths(
        run_dir, run_dir / "checkpoints", run_dir / "samples", run_dir / "metrics"
    )
    return cfg, paths

check_compatibility ¶

check_compatibility(model: LightningModule, datamodule: ECGDataModule) -> None

Fail early if the model and dataset disagree on leads, length, or conditions.

Source code in src/ecgen/training/runner.py

def check_compatibility(model: pl.LightningModule, datamodule: ECGDataModule) -> None:
    """Fail early if the model and dataset disagree on leads, length, or conditions."""
    ds = datamodule.reference_dataset
    conditioner = getattr(model, "conditioner", None)
    if conditioner is not None:
        sample = ds[0]
        available = set(sample.get("cond", {}))
        missing = [f for f in conditioner.cond_fields if f not in available]
        if missing:
            raise ValueError(
                f"Conditioning uses fields {missing} that the dataset does not provide "
                f"(it provides {sorted(available)})"
            )
        if conditioner.uses_obs and datamodule.mask_spec is None:
            raise ValueError(
                "Conditioning reads the observed ECG ('obs' field) but data.masks is not set"
            )
    io = _model_io(model)
    if io["channels"] is not None and io["channels"] != ds.num_leads:
        raise ValueError(
            f"Model expects {io['channels']} channels but the dataset yields {ds.num_leads} "
            f"leads {ds.lead_names}; set data.dataset.params.leads or the model's channel count"
        )
    if io["seq_length"] is not None and io["seq_length"] != ds.seq_length:
        raise ValueError(
            f"Model expects seq_length={io['seq_length']} but the dataset yields {ds.seq_length}"
        )

train ¶

train(cfg: Mapping[str, Any], overrides: Sequence[str] = (), config_path: Optional[str | Path] = None) -> pl.Trainer

Train a model from a config dict; returns the fitted pl.Trainer.

Parameters:

Name	Type	Description	Default
`cfg`	`Mapping[str, Any]`	Parsed experiment config.	required
`overrides`	`Sequence[str]`	`dotted.key=value` strings applied on top of `cfg`.	`()`
`config_path`	`Optional[str \| Path]`	Source YAML, recorded in the run metadata.	`None`

Source code in src/ecgen/training/runner.py

def train(
    cfg: Mapping[str, Any],
    overrides: Sequence[str] = (),
    config_path: Optional[str | Path] = None,
) -> pl.Trainer:
    """
    Train a model from a config dict; returns the fitted ``pl.Trainer``.

    Args:
        cfg: Parsed experiment config.
        overrides: ``dotted.key=value`` strings applied on top of ``cfg``.
        config_path: Source YAML, recorded in the run metadata.
    """
    cfg, paths = resolve_config(apply_overrides(dict(cfg), overrides))
    seed = cfg["experiment"]["seed"]
    set_global_seed(seed)

    model = build_model(cfg["model"])
    datamodule = build_datamodule(cfg["data"])
    datamodule.setup("fit")
    check_compatibility(model, datamodule)
    if hasattr(model, "set_data_info"):
        model.set_data_info(datamodule.state_dict())

    for d in (paths.checkpoints, paths.samples, paths.metrics):
        d.mkdir(parents=True, exist_ok=True)
    write_yaml(cfg, paths.run_dir / "config_resolved.yaml")
    write_json(
        collect_run_metadata(
            config_path=Path(config_path) if config_path else None,
            cfg=cfg,
            run_dir=paths.run_dir,
            argv=None,
        ),
        paths.run_dir / "metadata.json",
    )

    checkpoint = ModelCheckpoint(
        dirpath=str(paths.checkpoints),
        auto_insert_metric_name=False,
        **cfg["checkpoint"],
    )
    callbacks = [checkpoint] + build_callbacks(cfg["callbacks"], paths.samples)

    trainer = pl.Trainer(
        default_root_dir=str(paths.run_dir),
        logger=_loggers(cfg, paths.run_dir),
        callbacks=callbacks,
        **cfg["trainer"],
    )
    logger.info(
        f"Training '{cfg['experiment']['name']}' ({type(model).__name__}) in {paths.run_dir}"
    )
    trainer.fit(model=model, datamodule=datamodule, ckpt_path=cfg["resume"])

    write_json(
        {
            "callback_metrics": {
                k: float(v) for k, v in trainer.callback_metrics.items()
            },
            "logged_metrics": {k: float(v) for k, v in trainer.logged_metrics.items()},
            "best_model_path": checkpoint.best_model_path,
        },
        paths.metrics / "train_val_metrics.json",
    )
    logger.info(f"Done. Best checkpoint: {checkpoint.best_model_path or '<none>'}")
    return trainer

`ecgen.cli`¶

Command-line interface:

ecgen train <config.yaml> [key=value ...]
ecgen sample <checkpoint> -n 16 -o out --format csv [--cond field=value ...]
ecgen evaluate <samples.npy> <config.yaml>
ecgen list

parse_cond ¶

parse_cond(items: Sequence[str]) -> dict

["a=1,2,nan", "report=Sinus rhythm"] -> {"a": [1.0, 2.0, nan], "report": "Sinus rhythm"}.

Source code in src/ecgen/cli.py

def parse_cond(items: Sequence[str]) -> dict:
    """``["a=1,2,nan", "report=Sinus rhythm"]`` -> ``{"a": [1.0, 2.0, nan], "report": "Sinus rhythm"}``."""
    out = {}
    for item in items:
        if "=" not in item:
            raise ValueError(f"--cond must look like FIELD=VALUE, got {item!r}")
        key, raw = item.split("=", 1)
        raw = raw.strip().strip("'\"")
        parts = [x.strip() for x in raw.split(",")]
        try:
            values = [float(x) for x in parts]
            out[key] = values if len(values) > 1 else values[0]
        except ValueError:
            out[key] = raw
    return out

`ecgen.registry`¶

Name → class registries used to build components from YAML.

A component spec is a mapping with either a registry name or an import path target, plus optional params::

{"name": "vae", "params": {...}}
{"target": "my_pkg.models.MyModel", "params": {...}}

target lets user code outside the library plug in without registering.

Registry ¶

Bases: Generic[T]

A named collection of classes (or factories) of one kind.

Example::

FAMILIES = Registry("family")

@FAMILIES.register("vae")
class VAELightning(...): ...

FAMILIES.get("vae")  # -> VAELightning