Посібник із програмування Kauldron від Google Research: конфігурації як звичайні дані, компоненти, пов’язані рядками, і тренер JAX, який можна прочитати від початку до кінця

import os
import sys
import json
import textwrap
import traceback
import subprocess
 
RESULTS = {}
 
 
def banner(title):
    print("\n" + "=" * 78)
    print(title)
    print("=" * 78)
 
 
def section(name):
    def wrap(fn):
        def run(*a, **kw):
            banner(name)
            try:
                out = fn(*a, **kw)
                RESULTS[name] = out if isinstance(out, str) else "ok"
                return out
            except Exception as e:
                RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"
                print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")
                traceback.print_exc(limit=3)
                return None
        return run
    return wrap
 
 
banner("0. Install Kauldron, and the one compatibility patch you need today")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "kauldron==1.4.2"], check=True)
 
import jax
from etils.enp import array_spec as _array_spec
 
# jax >= 0.10.1 moved `jax._src.prng`, but etils <= 1.14.0 still reaches for it whenever it
# inspects an array's dtype. Kauldron calls that code on every batch, so without this two-line
# patch a Trainer raises AttributeError before it finishes a single step. The replacement uses
# jax's own public dtype API and is a no-op on older jax.
if not hasattr(jax._src, "prng"):
    _array_spec._is_jax_random_dtype = lambda dt: jax.dtypes.issubdtype(dt, jax.dtypes.prng_key)
 
import numpy as np
import optax
import flax
from flax import linen as nn
import kauldron
from kauldron import kd, konfig, kontext
from kauldron.typing import Float, typechecked
 
print(f"  kauldron {kauldron.__version__}  |  jax {jax.__version__}  |  flax {flax.__version__}"
      f"  |  optax {optax.__version__}")
print(f"  devices: {jax.devices()}")
print("\n  Kauldron's pitch is modularity: it is the glue, not the framework. Four pieces do the work:")
print("    konfig    -> your experiment IS a Python call tree, and that tree is a plain dict")
print("    kontext   -> parts are wired by string key paths, so they never import each other")
print("    ktyping   -> Float['*b h w c'] checked at runtime, with named axes bound across args")
print("    kd.train  -> Trainer: model + data + losses + metrics + optimizer, and nothing else")
print("\n  Everything below runs on a CPU runtime with no dataset download: the data is synthetic.")

@section("1. A config is a call tree, and a call tree is a dict")
def config_is_a_dict():
    with konfig.imports():
        import optax as coptax                      # looks like optax, builds ConfigDict instead
 
    cfg = coptax.adam(learning_rate=0.003)
    print(f"  cfg          = {cfg}")
    print(f"  type         = {type(cfg).__name__}")
    print(f"  __qualname__ = {cfg.__qualname__!r}   <- the call, stored as data")
 
    cfg.learning_rate = 1e-4                        # configs are mutable
    optimizer = konfig.resolve(cfg)                 # ...until you resolve them
    print(f"  after cfg.learning_rate = 1e-4 -> resolve() gives {type(optimizer).__name__}")
 
    print("\n  An arbitrarily complex optimizer is still just nested dicts:")
    chain = coptax.chain(
        coptax.clip_by_global_norm(1.0),
        coptax.scale_by_adam(b2=0.99),
        coptax.scale_by_learning_rate(0.003),
    )
    as_json = json.dumps(json.loads(chain.to_json()), indent=2)
    print(textwrap.indent(as_json, "    "))
 
    rebuilt = konfig.resolve(konfig.ConfigDict(json.loads(chain.to_json())))
    print(f"  JSON -> ConfigDict -> resolve() -> {type(rebuilt).__name__}")
    print("  optax has no idea konfig exists. No base class, no registry, no decorator.")
    return f"optax.chain -> JSON -> {type(rebuilt).__name__}"
 
 
config_is_a_dict()

@section("2. cfg.ref: change one number, everything downstream follows")
def config_references():
    with konfig.imports():
        import optax as coptax
        from kauldron import kd as ckd
 
    cfg = ckd.train.Trainer()
    cfg.num_train_steps = 1000
    cfg.schedules = {
        "lr": coptax.warmup_cosine_decay_schedule(
            init_value=0.0, peak_value=1e-3, warmup_steps=100,
            decay_steps=cfg.ref.num_train_steps,    # <- a reference, not the value 1000
        )
    }
    at_1000 = konfig.resolve(cfg.schedules["lr"])
    cfg.num_train_steps = 200                       # one edit...
    at_200 = konfig.resolve(cfg.schedules["lr"])    # ...and the schedule already knows
 
    print(f"  {'progress':>10s}  {'lr @ 1000 steps':>18s}  {'lr @ 200 steps':>16s}")
    for frac in (0.1, 0.5, 0.9):
        print(f"  {frac:>9.0%}  {float(at_1000(int(1000*frac))):>18.6f}"
              f"  {float(at_200(int(200*frac))):>16.6f}")
    print("\n  Without .ref the schedule would have frozen 1000 into itself, and a sweep over")
    print("  num_train_steps would have silently trained on the wrong decay curve.")
    return (f"lr at 90% of training: {float(at_1000(900)):.6f} (1000 steps)"
            f" vs {float(at_200(180)):.6f} (200 steps)")
 
 
config_references()

@section("3. kontext: parts are wired by string, so they never import each other")
def kontext_keys():
    import dataclasses
 
    ctx = {
        "batch": {"image": np.zeros((4, 8, 8, 3)), "label": np.arange(4)},
        "preds": {"logits": np.ones((4, 10)), "aux": [{"pos": np.zeros(3)}]},
    }
    print("  a context is just nested data; a key path reaches into it:")
    for path in ["batch.image", "preds.logits", "preds.aux[0].pos"]:
        print(f"    {path:22s} -> {kontext.get_by_path(ctx, path).shape}")
 
    try:
        kontext.get_by_path(ctx, "batch.nope")
    except KeyError as e:
        print(f"    {'batch.nope':22s} -> KeyError: {str(e)[:96]}...")
 
    @dataclasses.dataclass(eq=True, frozen=True, kw_only=True)
    class MeanGap:
        preds: kontext.Key = kontext.REQUIRED       # these ARE the wiring
        targets: kontext.Key = kontext.REQUIRED
 
        def __call__(self, *, preds, targets):
            return float(abs(np.asarray(preds).mean() - np.asarray(targets).mean()))
 
    metric = MeanGap(preds="preds.logits", targets="batch.label")
    kwargs = kontext.resolve_from_keyed_obj(ctx, metric)
    print(f"\n  MeanGap declared preds={metric.preds!r}, targets={metric.targets!r}")
    print(f"  resolved to kwargs: {{{', '.join(f'{k}: {v.shape}' for k, v in kwargs.items())}}}")
    print(f"  value = {metric(**kwargs)}")
    print("\n  MeanGap never imported the model and the model never heard of MeanGap. Point the")
    print("  same metric at 'preds.aux[0].pos' and nothing but that string changes.")
    return f"MeanGap(preds='preds.logits', targets='batch.label') = {metric(**kwargs)}"
 
 
kontext_keys()

@section("4. ktyping: named axes, checked at runtime, bound across arguments")
def shape_checking():
    @typechecked
    def project(features: Float["*b n c"], weights: Float["c d"]) -> Float["*b n d"]:
        return jax.numpy.einsum("...c,cd->...d", features, weights)
 
    out = project(jax.numpy.zeros((2, 16, 8)), jax.numpy.zeros((8, 32)))
    print(f"  project(f32[2 16 8], f32[8 32]) -> {out.shape}   c bound to 8, d bound to 32")
 
    print("\n  now break it: c is bound to 8 by the first argument, so 5 cannot also be c")
    try:
        project(jax.numpy.zeros((2, 16, 8)), jax.numpy.zeros((5, 32)))
    except Exception as e:
        print(textwrap.indent(str(e), "    "))
    print("\n  'Inferred Dims' is the part worth having: it reports what each axis name was already")
    print("  bound to, so a mismatch names the axis instead of printing two anonymous shapes.")
    return "mismatch named the axis: c already bound to 8, got 5"
 
 
shape_checking()

import dataclasses
 
 
@dataclasses.dataclass(eq=True, frozen=True, kw_only=True)
class LogCosh(kd.losses.Loss):
    """log(cosh(err)): quadratic near zero, linear in the tails. ~30 lines less than raw Flax."""
 
    preds: kontext.Key = kontext.REQUIRED
    targets: kontext.Key = kontext.REQUIRED
 
    @typechecked
    def get_values(self, preds: Float["*a"], targets: Float["*a"]) -> Float["*a"]:
        return jax.numpy.log(jax.numpy.cosh(preds - targets))
 
 
@dataclasses.dataclass(eq=True, frozen=True, kw_only=True)
class WithinTol(kd.metrics.Metric):
    """Fraction of predictions landing within `tol` of the target, over every batch seen."""
 
    preds: kontext.Key = kontext.REQUIRED
    targets: kontext.Key = kontext.REQUIRED
    tol: float = 0.25
 
    @flax.struct.dataclass
    class State(kd.metrics.AutoState):
        # sum_field() marks a value that is ADDED when two states merge. Keeping the numerator
        # and the denominator apart is what makes the pooled result exact.
        n_hit: Float[""] = kd.metrics.sum_field(default=0.0)
        n_total: Float[""] = kd.metrics.sum_field(default=0.0)
 
        def compute(self) -> Float[""]:
            # Return a jax scalar, like the built-in states do: the metric writer that
            # `trainer.train()` logs through does not accept a bare numpy scalar.
            total = jax.numpy.maximum(jax.numpy.asarray(self.n_total), 1.0)
            return jax.numpy.asarray(self.n_hit) / total
 
    @typechecked
    def get_state(self, preds: Float["*a"], targets: Float["*a"]) -> "WithinTol.State":
        hit = (jax.numpy.abs(preds - targets) < self.tol).astype("float32")
        return self.State(n_hit=hit.sum(), n_total=jax.numpy.asarray(hit.size, "float32"))
 
 
@section("5. A custom loss and a custom metric, in the shape Kauldron expects")
def custom_loss_and_metric():
    rng = np.random.default_rng(0)
    p = jax.numpy.asarray(rng.normal(size=(8, 4)).astype("float32"))
    t = jax.numpy.asarray(rng.normal(size=(8, 4)).astype("float32"))
 
    loss = LogCosh(preds="preds.y", targets="batch.y")
    print(f"  {'LogCosh(preds, targets)':34s} {float(loss(preds=p, targets=t)):.6f}")
    print(f"  {'same loss, weight=0.5':34s} "
          f"{float(LogCosh(preds='a', targets='b', weight=0.5)(preds=p, targets=t)):.6f}   <- exactly half")
    print(f"  {'builtin kd.losses.L2':34s} {float(kd.losses.L2(preds='a', targets='b')(preds=p, targets=t)):.6f}")
 
    print("\n  A metric is not a number, it is a State that merges. Watch why that matters when")
    print("  the last batch of an epoch is smaller than the rest:")
    metric = WithinTol(preds="preds.y", targets="batch.y", tol=0.5)
    big = metric.get_state(preds=p[:6], targets=t[:6])
    small = metric.get_state(preds=p[6:], targets=t[6:])
    merged = big.merge(small)
    for label, st in [("batch of 6 rows", big), ("batch of 2 rows", small), ("big.merge(small)", merged)]:
        print(f"    {label:22s} {float(st.n_hit):>4.0f} / {float(st.n_total):>3.0f}  = {float(st.compute()):.4f}")
    naive = (float(big.compute()) + float(small.compute())) / 2
    print(f"    {'mean of the two rates':22s} {'':>4s}   {'':>3s}  = {naive:.4f}   <- wrong, and quietly so")
    print("\n  sum_field() adds numerator and denominator separately, so the pooled value is exact")
    print("  however the batches were sized. The same mechanism aggregates a metric across devices:")
    print("  merge is associative, so the order the states arrive in never changes the answer.")
    print("\n  Subclass, annotate the keys, implement one method. The loss never sees a batch dict")
    print("  and the metric never sees the model; the keys deliver exactly what was asked for.")
    return (f"merged {float(merged.n_hit):.0f}/{float(merged.n_total):.0f} = {float(merged.compute()):.4f}"
            f" vs {naive:.4f} from averaging the rates")
 
 
custom_loss_and_metric()

SEED_W = np.random.default_rng(7).normal(size=(16, 1)).astype("float32")
 
 
def make_loader(split: str):
    """A dataset is any callable returning an array tree. No download, no tf.data, no TFDS."""
 
    def load():
        rng = np.random.default_rng(0 if split == "train" else 1)
        n = 512 if split == "train" else 128
        x = rng.normal(size=(n, 16)).astype("float32")
        y = (x @ SEED_W + 0.1 * rng.normal(size=(n, 1))).astype("float32")
        return {"x": x, "y": y}
 
    return load
 
 
class MLP(nn.Module):
    inputs: kontext.Key = kontext.REQUIRED          # the model declares where its input comes from
    hidden: int = 32
 
    @nn.compact
    def __call__(self, inputs: Float["*b f"]) -> dict[str, Float["*b 1"]]:
        h = nn.relu(nn.Dense(self.hidden, name="enc")(inputs))
        return {"y": nn.Dense(1, name="out")(h)}
 
 
@section("6. A real Trainer on synthetic data, and a look inside the model")
def train_for_real():
    train_ds = kd.data.InMemoryPipeline(
        loader=make_loader("train"), batch_size=32, shuffle=True, num_epochs=None, seed=0)
    print(f"  element_spec: {kd.inspect.json_spec_like(train_ds.element_spec)}")
    print("\n  batch statistics straight from kd.inspect:")
    print(textwrap.indent(str(kd.inspect.get_batch_stats(next(iter(train_ds)))), "    "))
 
    trainer = kd.train.Trainer(
        seed=0,
        workdir="/tmp/kauldron_tutorial",
        train_ds=train_ds,
        model=MLP(inputs="batch.x"),
        num_train_steps=300,
        train_losses={"logcosh": LogCosh(preds="preds.y", targets="batch.y")},
        train_metrics={
            "within_tol": WithinTol(preds="preds.y", targets="batch.y"),
            "enc_norm": kd.metrics.Norm(tensor="interms.enc.__call__[0]"),   # an INNER layer
        },
        optimizer=optax.adam(learning_rate=1e-2),
    )
 
    it = iter(trainer.train_ds)
    state = trainer.trainstep.init(trainer.train_ds.element_spec)
    print(f"\n  params: {jax.tree.map(lambda x: tuple(x.shape), state.params)}")
    print(f"\n  {'step':>5s} {'logcosh':>10s} {'within .25':>11s} {'enc_norm':>10s}")
    for step in range(1, 301):
        # aux is opt-in: the step skips building it unless you ask, because it costs device time.
        state, aux = trainer.trainstep.step(state, next(it), return_losses=True, return_metrics=True)
        if step in (1, 25, 50, 100, 200, 300):
            losses = {k: float(v.compute()) for k, v in aux.loss_states.items()}
            metrics = {k: float(v.compute()) for k, v in aux.metric_states.items()}
            first_loss = losses["logcosh"] if step == 1 else first_loss
            final_loss = losses["logcosh"]
            print(f"  {step:>5d} {losses['logcosh']:>10.4f} {metrics['within_tol']:>11.4f}"
                  f" {metrics['enc_norm']:>10.4f}")
 
    print("\n  enc_norm was never returned by the model. 'interms.enc.__call__[0]' reaches into the")
    print("  Dense layer named 'enc' through Flax's captured intermediates, so monitoring an inner")
    print("  activation costs one string in the config and zero edits to MLP.")
    return f"logcosh {first_loss:.4f} -> {final_loss:.4f} over 300 CPU steps"
 
 
train_for_real()

@section("7. A sweep is a for-loop over config overrides")
def sweep():
    with konfig.imports():
        import optax as coptax
        from kauldron import kd as ckd
    with konfig.imports(lazy=True):
        # Classes defined in a notebook live in __main__, which cannot be fake-imported eagerly.
        from __main__ import MLP as CfgMLP
        from __main__ import make_loader as cfg_make_loader
        from __main__ import LogCosh as CfgLogCosh
 
    def base_config():
        cfg = ckd.train.Trainer()
        cfg.seed = 0
        cfg.workdir = "/tmp/kauldron_sweep"
        cfg.train_ds = ckd.data.InMemoryPipeline(
            loader=cfg_make_loader("train"), batch_size=32, shuffle=True, num_epochs=None)
        cfg.model = CfgMLP(inputs="batch.x", hidden=32)
        cfg.num_train_steps = 200
        cfg.train_losses = {"logcosh": CfgLogCosh(preds="preds.y", targets="batch.y")}
        cfg.optimizer = coptax.adam(learning_rate=1e-2)
        return cfg
 
    print(f"  cfg.model     = {base_config().model}")
    print(f"  cfg.optimizer = {base_config().optimizer}")
    print("\n  A bare lazy-imported name is a reference, not a call:")
    print(f"    CfgMLP(inputs=...)  -> {{'__qualname__': '__main__.MLP', ...}}   (built when resolved)")
    print(f"    cfg_make_loader     -> {dict(cfg_make_loader)}   (handed over as-is)")
 
    print(f"\n  {'override':34s} {'final logcosh':>14s}")
    results = {}
    for label, apply_override in [
        ("(baseline)", lambda c: None),
        ("cfg.model.hidden = 4", lambda c: setattr(c.model, "hidden", 4)),
        ("cfg.model.hidden = 128", lambda c: setattr(c.model, "hidden", 128)),
        ("cfg.optimizer.learning_rate = 0.1", lambda c: setattr(c.optimizer, "learning_rate", 0.1)),
        ("cfg.optimizer = optax.sgd(0.05)", lambda c: setattr(c, "optimizer", coptax.sgd(0.05))),
    ]:
        cfg = base_config()
        apply_override(cfg)
        trainer = konfig.resolve(cfg)               # ConfigDict -> a real, frozen Trainer
        it = iter(trainer.train_ds)
        state = trainer.trainstep.init(trainer.train_ds.element_spec)
        for _ in range(200):
            state, aux = trainer.trainstep.step(state, next(it), return_losses=True)
        results[label] = float(aux.loss_states["logcosh"].compute())
        print(f"  {label:34s} {results[label]:>14.4f}")
 
    print("\n  Five experiments, five one-line edits, and not one character of MLP, LogCosh or the")
    print("  training loop changed. On the command line the same overrides are")
    print("  --cfg.model.hidden=128, which is why a Kauldron sweep is a list of these strings.")
    return (f"best {min(results.values()):.4f} ({min(results, key=results.get)}),"
            f" worst {max(results.values()):.4f} ({max(results, key=results.get)})")
 
 
sweep()

@section("8. The guardrail: configs hold configs, never resolved objects")
def guardrails():
    with konfig.imports():
        from kauldron import kd as ckd
 
    cfg = ckd.train.Trainer()
    print("  Assigning a REAL flax module into a config is refused on the spot:")
    try:
        cfg.model = MLP(inputs="batch.x", hidden=8)
    except ValueError as e:
        print(textwrap.indent(str(e)[:420], "    "))
 
    print("\n  Why this matters: a half-resolved config cannot be serialized, diffed or overridden")
    print("  from the command line, so konfig refuses to let one exist rather than failing later.")
 
    print("\n  The same discipline shows up in sub-objects, which default to root-config references:")
    print(f"    kd.data.InMemoryPipeline(...).seed   default -> _FakeRootCfg('cfg.seed')")
    print(f"    kd.evals.Evaluator(...).ds           default -> _FakeRootCfg('cfg.eval_ds')")
    print("  Inside a Trainer those are filled from the root. Built standalone they are not, which")
    print("  is why step 6 passed seed=0 to the pipeline explicitly.")
    return "ConfigDict refused a resolved flax module at assignment"
 
 
guardrails()

@section("9. Evaluation and checkpointing, and picking up where you left off")
def eval_and_checkpoint():
    import pathlib
    import shutil
 
    workdir = "/tmp/kauldron_resume"
    shutil.rmtree(workdir, ignore_errors=True)      # start from a clean slate for the demo
 
    def build(num_steps):
        return kd.train.Trainer(
            seed=0,
            workdir=workdir,
            train_ds=kd.data.InMemoryPipeline(
                loader=make_loader("train"), batch_size=32, shuffle=True, num_epochs=None, seed=0),
            model=MLP(inputs="batch.x"),
            num_train_steps=num_steps,
            log_metrics_every=100,
            train_losses={"logcosh": LogCosh(preds="preds.y", targets="batch.y")},
            train_metrics={"within_tol": WithinTol(preds="preds.y", targets="batch.y")},
            optimizer=optax.adam(learning_rate=1e-2),
            checkpointer=kd.ckpts.Checkpointer(save_interval_steps=100),
            evals={
                "eval": kd.evals.Evaluator(
                    run=kd.evals.EveryNSteps(100),
                    ds=kd.data.InMemoryPipeline(
                        loader=make_loader("eval"), batch_size=32, shuffle=False,
                        num_epochs=1, seed=0),
                    num_batches=4,
                )
            },
        )
 
    state, _ = build(200).train()
    print(f"  first run finished at step {int(state.step)}")
    saved = sorted(p.name for p in pathlib.Path(workdir).glob("checkpoints/ckpt_*"))
    print(f"  checkpoints on disk: {saved}")
 
    print("\n  Now build the same Trainer again, on the same workdir, asking for more steps:")
    state2, _ = build(300).train()
    print(f"  second run finished at step {int(state2.step)}")
    print("  The progress bar above started at 200, not 0: train() found the checkpoint and")
    print("  resumed, which is also what happens when a preemptible job is restarted.")
    print("\n  The evaluator ran on its own dataset every 100 steps. It inherited the model, the")
    print("  losses and the metrics from the root config, so declaring it took four lines.")
    return f"trained, evaluated, checkpointed and resumed at step {int(state2.step)}"
 
 
eval_and_checkpoint()

banner("SUMMARY")
for name, res in RESULTS.items():
    print(f"  {name:<72s}  {res}")
print("""
Where to go next
 - Read the two files that carry the ideas: kauldron/konfig/ (the config system, usable in any
   project with `from kauldron import konfig`) and kauldron/kontext/ (the key system). Both are
   self-contained and have no dependency on the rest of Kauldron.
 - Swap the synthetic pipeline for a real one: kd.data supports TFDS, Grain and PyGrain sources,
   plus transforms (Resize, Rearrange, ValueRange, Elements) that are themselves config entries.
 - Start from a working config: github.com/google-research/kauldron/tree/main/examples
 - Run it as a script: `python -m kauldron.main --cfg=config.py --cfg.model.hidden=128`, which is
   the same override used in step 7, typed on the command line instead.
 - Docs: kauldron.readthedocs.io  (konfig, kontext, data, eval, sharding, checkpoints)
""")