File size: 3,251 Bytes
534cc76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """W&B metrics, checkpoint artifacts and run continuation."""
from contextlib import contextmanager
from .checkpoint import save_checkpoint
from .provenance import VERSION
@contextmanager
def open_run(config, directory, model, datasets, previous=None):
"""Create a tracking context; disabled mode does not initialize W&B."""
settings = config.get("wandb", {})
mode = settings.get("mode", "disabled")
interval = settings.get("log_every", 10)
if mode not in ("online", "offline", "disabled"):
raise ValueError("wandb.mode must be online, offline, or disabled")
if isinstance(interval, bool) or not isinstance(interval, int) or interval < 1:
raise ValueError("wandb.log_every must be a positive integer")
if mode == "disabled":
yield None
return
import wandb
project = settings.get("project", "paya")
entity = settings.get("entity")
continuation = {}
if previous and mode == previous["mode"] == "online":
# W&B can return canonical capitalization different from the requested project.
if project.casefold() != previous["project"].casefold() or entity not in (None, previous["entity"]):
raise ValueError("Online resume must use the saved W&B project and entity")
project = previous["project"]
entity = previous["entity"]
continuation = {"id": previous["id"], "resume": "must"}
with wandb.init(project=project, entity=entity, name=settings.get("name") or directory.name,
mode=mode, dir=str(directory), save_code=False, **continuation,
settings=wandb.Settings(init_timeout=90,finish_timeout=settings.get("finish_timeout",120),
finish_timeout_raises=True),
config=None if continuation else {"version": VERSION, "model": model.identity(),
"training": {k: v for k, v in config.items()
if k not in ("wandb", "data", "assets", "feature_cache")},
"datasets": {name: {"records": len(data.rows), "sha256": data.sha256}
for name, data in datasets.items()},
"resumed_from_run_id": previous["id"] if previous else None}) as run:
run.define_metric("step")
run.define_metric("epoch")
run.define_metric("train/*", step_metric="step")
run.define_metric("selection/*", step_metric="epoch")
yield run
def run_identity(run, mode):
return {"id": run.id, "project": run.project, "entity": run.entity, "mode": mode} if run else None
def log_best(run, model, directory, epoch, nll):
"""Log inference weights and their selection result, excluding optimizer/data state."""
import wandb
path = directory/"best-decision.pt"
save_checkpoint(path, model)
artifact = wandb.Artifact(f"paya-{VERSION}-{run.id}", type="model",
metadata={"epoch": epoch, "selection_nll": nll, "version": VERSION})
artifact.add_file(str(path), name="decision.pt")
run.log_artifact(artifact, aliases=["best"])
run.summary.update({"best/epoch": epoch, "best/selection_nll": nll})
|