在本教程中，我们实现 Kauldron，这是来自 Google Research 的 JAX 训练库，它自称针对研究速度和模块化进行了优化，而我们按字面理解这两个词，测试它们实际能带来什么。我们先安装它，然后用笔记本的前半部分讲解让 Kauldron 区别于一堆 Flax 和 Optax 的三个机制：konfig，它把实验变成一棵可以通过 JSON 往返转换的普通字典树；kontext，它用字符串键路径把组件连接起来，使得损失函数永远不需要导入它所评分的模型；以及运行时形状检查器，它的命名轴会跨参数绑定，并在不匹配时报告它们被绑定到了什么。然后我们按照框架期望的形式编写自定义损失和自定义指标，在合成的内存数据上训练一个真正的 Trainer，无需下载也无需加速器，并在不修改模型的情况下监控模型的内部层。最后我们运行一个五变体扫描，其中每个实验只相差一行配置，并让一次训练运行自行检查点并在停止处恢复。
复制代码已复制使用其他浏览器
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. 安装 Kauldron，以及你今天需要的一个兼容性补丁")
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 移动了 `jax._src.prng`，但 etils <= 1.14.0 在检查数组的 dtype 时
# 仍然会去访问它。Kauldron 在每个批次上都会调用那段代码，所以没有这个两行
# 补丁，Trainer 在完成一步之前就会抛出 AttributeError。替换方案使用
# jax 自己的公开 dtype API，在旧版 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 的卖点是模块化：它是胶水，而不是框架。四部分承担了工作：")
print(" konfig -> 你的实验就是一个 Python 调用树，而这棵树就是一个普通 dict")
print(" kontext -> 各部分通过字符串键路径连接，因此它们从不互相 import")
print(" ktyping -> Float['*b h w c'] 在运行时检查，命名轴跨参数绑定")
print(" kd.train -> Trainer：模型 + 数据 + 损失 + 指标 + 优化器，别无其他")
print("\n 下面所有内容都在 CPU 运行时上运行，无需下载数据集：数据是合成的。")
我们安装 Kauldron，并应用当前版本组合所需的那个兼容性补丁。jax 0.10.1 移动了私有模块 jax._src.prng，而 etils 直到 1.14.0 在检查数组的 dtype 时仍会去访问它，这是 Kauldron 在每个批次上都会运行的代码路径。如果没有下面这两行替换代码（它使用 jax 自己的公开 dtype API，在旧版本上是空操作），Trainer 会在完成一步之前就抛出 AttributeError。有了它，我们导入承担工作的四个部分：用于配置系统的 konfig、用于连接各部分的 kontext、用于运行时形状检查的 typing 模块，以及用于 Trainer 本身的 kd.train。之后的一切都在 CPU 运行时上运行，因为这个 notebook 中唯一的数据集是我们自己生成的。
复制代码已复制使用其他浏览器
@section("1. 配置就是一棵调用树，而调用树就是一个 dict")
def config_is_a_dict():
with konfig.imports():
import optax as coptax # 看起来像 optax，实际构建的是 ConfigDict
cfg = coptax.adam(learning_rate=0.003)
print(f" cfg = {cfg}")
print(f" type = {type(cfg).__name__}")
print(f" __qualname__ = {cfg.__qualname__!r} <- 这个调用，作为数据存储")
cfg.learning_rate = 1e-4 # 配置是可变的
optimizer = konfig.resolve(cfg) # ...直到你解析它们
print(f" 在 cfg.learning_rate = 1e-4 之后 -> resolve() 得到 {type(optimizer).__name__}")
print("\n 任意复杂的优化器仍然只是嵌套字典：")
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 完全不知道 konfig 的存在。没有基类，没有注册表，没有装饰器。")
return f"optax.chain -> JSON -> {type(rebuilt).__name__}"
config_is_a_dict()
我们从 konfig 开始，因为它是整个库其余部分所构建的基础。在 konfig.imports() 块内，导入 optax 会给我们一个看起来和自动补全都像 optax 的东西，但它构建的是配置而不是对象，因此 optax.adam(learning_rate=0.003) 返回的是一个 ConfigDict，其中保存着该调用的限定名称及其参数，而不是一个优化器。该配置在 konfig.resolve 将其变成真正的对象之前是可变的，而且因为它只是嵌套字典，任意复杂的 optax.chain 都能序列化为 JSON，并重新变回一个可用的优化器。重要的部分是 optax 为了支持这一点所必须做的事情：什么都没有。optax 中没有任何基类、注册表或装饰器，我们以这种方式配置的任何库也同样如此。
复制代码已复制使用其他浏览器
@section("2. cfg.ref：改一个数字，下游的一切都随之改变")
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, # <- 一个引用，而不是值 1000
)
}
at_1000 = konfig.resolve(cfg.schedules["lr"])
cfg.num_train_steps = 200 # 一处修改……
at_200 = konfig.resolve(cfg.schedules["lr"]) # ……调度器已经知道了
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()
配置系统通常会在同一个值需要在多个地方使用时出错，而 Kauldron 的答案是 cfg.ref。我们把 warmup-cosine 调度器的 decay_steps 指向 cfg.ref.num_train_steps 而不是 1000，然后把 num_train_steps 改成 200 并重新解析该调度器。学习率曲线会自行重塑，因为配置存储的是引用而不是值的副本。如果没有这层间接引用，调度器就会把 1000 冻结进自身，而对训练步数做扫描时，每个变体都会悄悄地在错误的衰减曲线上训练，这正是那种会产生一个看似合理的数字却没有任何报错的 bug。
复制代码已复制使用其他浏览器
@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()
kontext 是 Kauldron 连接彼此毫无了解的组件的方式。context 就是普通的嵌套数据，而像 batch.image 或 preds.aux[0].pos 这样的键路径可以深入其中，解析字典键、属性和列表索引，并在无法解析时抛出一个 KeyError，列出实际可用的内容。任何对象都可以通过将字段标注为 kontext.Key 来声明其输入，然后 resolve_from_keyed_obj 会从 context 中精确地取出这些路径，并作为关键字参数交给它。我们以这种方式构建了一个小指标，并将其指向模型的输出：指标从不导入模型，模型也从未听说过该指标，而将指标重定向到另一个张量只需改动一个字符串。
Copy CodeCopiedUse a different Browser
@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()
Kauldron 的类型模块使用命名轴在运行时检查数组形状。我们用一个函数标注 Float['*b n c'] 和 Float['c d']，装饰器会在第一次看到每个轴名时将其绑定，然后在签名中的其他所有位置（包括返回值）强制执行该绑定。当我们故意传入一个不兼容的第二个参数时，错误信息做了真正重要的事：除了实际形状之外，它还打印出一个 Inferred Dims 块，显示 c 已经被绑定为 8，因此失败信息指出了不一致的那个轴，而不是让我们去比较两个匿名的元组。在一个同时有几十个张量在运行中的模型上，这就是一行修复和一次调试会话之间的区别。
Copy CodeCopiedUse a different Browser
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)
# `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. 自定义损失函数和自定义指标，采用 Kauldron 期望的形式")
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()
我们按照 Kauldron 期望的精确形式编写自定义损失和自定义指标，即一个带有 kontext.Key 字段和一个方法的冻结 dataclass。损失实现 get_values 并返回逐元素数组；框架处理归约和 weight 参数，我们通过检查 weight=0.5 是否恰好将结果减半来确认这一点。指标更有意思，因为 Kauldron 指标返回的不是一个数字，而是一个会合并的 State。我们基于 AutoState 构建自己的指标，包含两个 sum_field 条目，一个分子和一个分母，然后将一个六行的批次与一个两行的批次合并。汇总结果是精确的，而两个逐批次比率的平均值则明显错误，这恰恰会发生在某个 epoch 最后一批行数不齐时，并且由于 merge 具有结合性，同一机制可以在设备间聚合指标，而无需关心结果到达的顺序。
复制代码已复制使用其他浏览器
SEED_W = np.random.default_rng(7).normal(size=(16, 1)).astype("float32")
def make_loader(split: str):
"""数据集是任何返回数组树的 callable。无需下载，无需 tf.data，无需 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 # 模型声明其输入来自何处
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. 在合成数据上运行真正的 Trainer，并一窥模型内部")
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 直接来自 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]"), # 一个内部层
},
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 是可选启用的：除非你主动要求，否则该步骤会跳过构建它，因为它会消耗设备时间。
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 从未由模型返回。'interms.enc.__call__[0]' 通过 Flax 捕获的中间量深入到了",
print(" 名为 'enc' 的 Dense 层，因此监控一个内部激活只需在配置中写一个字符串，",
print(" 且对 MLP 零改动。")
return f"logcosh {first_loss:.4f} -> {final_loss:.4f} over 300 CPU steps"
train_for_real()
现在我们来训练。Kauldron 数据集是任何返回数组树的 callable，因此 kd.data.InMemoryPipeline 把我们的合成回归数据变成一个真正的流水线，具备批处理和打乱，且无需下载任何东西。我们用模型、该流水线、自定义损失、自定义指标和 Optax 优化器组装一个 Trainer，并直接驱动其 train step，以便打印损失曲线：在约 1 秒的 CPU 时间内，损失在三百步内从 1.51 降到 0.005。有两个细节值得记住。train step 不会构建其辅助输出，除非我们用 return_losses 和 return_metrics 请求，因为计算它们会消耗设备时间。而 enc_norm 列是从 interms.enc.__call__[0] 读取的，这是一个指向名为 enc 的 Dense 层中间输出的键路径，因此监控一个内部激活只需在配置中写一个字符串，完全不需要修改模型。
复制代码已复制使用其他浏览器
@section("7. 一次 sweep 就是对配置覆盖的 for 循环")
def sweep():
with konfig.imports():
import optax as coptax
from kauldron import kd as ckd
with konfig.imports(lazy=True):
# 在 notebook 中定义的类位于 __main__ 中，无法被急切地伪导入。
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 一个裸的惰性导入名称是引用，而不是调用：")
print(f" CfgMLP(inputs=...) -> {{'__qualname__': '__main__.MLP', ...}} (在解析时构建)")
print(f" cfg_make_loader -> {dict(cfg_make_loader)} (原样传递)")
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 -> 一个真实的、冻结的 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()
这一步就是整个设计的理由所在。我们把实验写一次作为配置，然后运行五个变体，每个变体恰好只差一行：两个模型宽度、两个优化器设置，以及把 Adam 整体换成 SGD。每个变体都解析成一个全新的 Trainer 并训练两百个真实步骤，而它们之间模型、损失或训练循环没有一个字符发生变化。有两个 konfig 细节让这在 notebook 中得以实现。在 notebook 中定义的类位于 __main__ 中，无法被急切地伪导入，因此我们对它们使用 konfig.imports(lazy=True)。而一个裸的惰性导入名称会解析为对象本身而不是调用它，这就是加载器函数被完整地交给 pipeline 的方式。在命令行中，这些相同的覆盖写作 –cfg.model.hidden=128，这就是为什么一个 Kauldron sweep 不过是这样一串字符串的列表。
复制代码已复制使用其他浏览器
@section("8. 护栏：配置持有配置，绝不持有已解析的对象")
def guardrails():
with konfig.imports():
from kauldron import kd as ckd
cfg = ckd.train.Trainer()
print(" 将真实的 flax 模块赋值到配置中会被当场拒绝：")
try:
cfg.model = MLP(inputs="batch.x", hidden=8)
except ValueError as e:
print(textwrap.indent(str(e)[:420], " "))
print("\n 为什么这很重要：一个半解析的配置无法被序列化、diff 或从命令行覆盖")
print(" ，所以 konfig 拒绝让这种状态存在，而不是留到之后才失败。")
print("\n 同样的纪律也体现在子对象上，它们默认引用根配置：")
print(f" kd.data.InMemoryPipeline(...).seed 默认 -> _FakeRootCfg('cfg.seed')")
print(f" kd.evals.Evaluator(...).ds 默认 -> _FakeRootCfg('cfg.eval_ds')")
print(" 在 Trainer 内部，这些字段会从根配置填充。单独构建时则不会，这")
print(" 就是为什么第 6 步显式地向 pipeline 传入了 seed=0。")
return "ConfigDict 在赋值时拒绝了已解析的 flax 模块"
guardrails()
Kauldron 拒绝让配置持有已解析的对象，而且这种拒绝是立即发生的，而非延迟的。将一个真实的 Flax 模块赋值到 ConfigDict 会当场抛出异常，并给出两条修复建议：把 import 包在 konfig.imports() 中，或在 notebook 中使用 mock_modules。一个半解析的配置无法被序列化、diff 或从命令行覆盖，因此系统直接排除这种状态，而不是留到之后以难以追踪的方式失败。同样的纪律也解释了我们在第 6 步遇到的情况：pipeline 和 evaluator 等子对象的 seed 和 dataset 字段默认引用根配置，当对象在 Trainer 内部构建时这些引用会被填充，而单独构建时则不会，这就是为什么我们向 pipeline 传入了显式的 seed。
Copy CodeCopiedUse a different Browser
@section("9. 评估与检查点，以及从上次中断处继续")
def eval_and_checkpoint():
import pathlib
import shutil
workdir = "/tmp/kauldron_resume"
shutil.rmtree(workdir, ignore_errors=True) # 为演示从干净状态开始
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()
我们用那些把训练脚本变成作业的部分来收尾。评估器只需四行代码就能声明，因为它从根配置继承了模型、损失和指标，只需要自己的数据集和一个说明运行频率的调度。检查点器按固定的步数间隔写入状态。然后我们针对同一个工作目录第二次构建同一个 Trainer，要求更多的步数，进度条从 200 而不是 0 开始：这次运行找到了它的检查点并恢复了。这与可抢占作业被重启时走的路径相同，值得在 notebook 里看一次——在那里整个过程只需一秒，而不是在集群上第一次发现它。
Copy CodeCopiedUse a different Browser
banner("SUMMARY")
for name, res in RESULTS.items():
print(f" {name:<72s} {res}")
print("""
接下来该往哪走
- 阅读承载这些理念的两个文件：kauldron/konfig/（配置系统，可在任何项目中通过 `from kauldron import konfig` 使用）以及 kauldron/kontext/（键系统）。两者都是
自包含的，不依赖 Kauldron 的其余部分。
- 把合成流水线换成真实的：kd.data 支持 TFDS、Grain 和 PyGrain 数据源，
外加各种变换（Resize、Rearrange、ValueRange、Elements），它们本身就是配置项。
- 从一个可用的配置开始：github.com/google-research/kauldron/tree/main/examples
- 把它当作脚本运行：`python -m kauldron.main --cfg=config.py --cfg.model.hidden=128`，这与
第 7 步中使用的覆盖方式相同，只不过是在命令行上输入的。
- 文档：kauldron.readthedocs.io （konfig、kontext、data、eval、sharding、checkpoints）
""")
摘要会打印每个部分返回的单行结果，然后指出接下来该往哪走：两个值得单独阅读的自包含模块、替代我们合成流水线的真实数据流水线、仓库中的示例配置，以及把第 7 步的扫描变成命令行标志的那条命令。
The summary prints the one-line result each section returned, then points at where to go next: the two self-contained modules worth reading on their own, the real data pipelines that replace our synthetic one, the example configs in the repository, and the command line that turns the sweep from step 7 into flags.
总之，我们把 Kauldron 的两个主张当作需要检验而非照搬的东西，而两者都成立，原因很简单。这里的模块化不是抽象层，而是没有抽象层：配置是一个描述 Python 调用的字典，连线是一个命名数据路径的字符串，optax 和我们自己的模型都不需要一行感知 Kauldron 的代码就能参与其中。研究速度由此而来，我们在 sweep 中具体看到了这一点——五个实验只花了五行编辑——以及在 shape checker 中，它指出不一致的轴，而不是打印两个形状然后让我们自己去琢磨。我们不需要的部分，TensorFlow datasets、Grain、sharding 和 XManager，完全让开了路，而 Trainer 在 CPU 上用合成数据运行。我们会带入实际工作的东西比这个库更小：把配置当作数据，通过路径而非 import 来连接组件，让指标成为可合并的状态，因为这三者在即使从不采用 Kauldron 本身的代码库中也很有用。
在这里查看完整代码。所有功劳归于这个项目的研究者。另外，欢迎在 Twitter 上关注我们，别忘了加入我们 150k+ 的 ML SubReddit 并订阅我们的 Newsletter。等等！你在 telegram 上吗？现在你也可以在 telegram 上加入我们了。
需要与我们合作来推广你的 GitHub Repo 或 Hugging Face Page 或产品发布或网络研讨会等？联系我们
文章 A Coding Guide to Google Research’s Kauldron: Configs That Are Plain Data, Components Wired by String, and a JAX Trainer You Can Read End to End 首次出现在 MarkTechPost 上。