kalfa 3.0.0.dev0__py3-none-any.whl

kalfa/std/loggers.py

@@ -0,0 +1,105 @@
+import sys
+
+from kalfa.core.aliases import register
+from kalfa.core.errors import BuildError
+from kalfa.core.protocols import Logger
+from kalfa.std.common import numeric, positive_int
+
+
+MISSING = "tensorboard logger needs the tensorboard package (uv sync installs the dependencies of kalfa)"
+COUNTERS = ("turn", "update", "seconds")
+PER_UPDATE = ("loss/", "lr/", "grad_norm/")
+
+
+def summary_writer():
+    try:
+        from torch.utils.tensorboard import SummaryWriter
+    except ImportError as exc:
+        raise BuildError(MISSING) from exc
+    return SummaryWriter
+
+
+def per_update(key):
+    return key.startswith(PER_UPDATE) and key.count("/") == 1
+
+
+def series_of(line):
+    return [key for key, value in line.items() if key not in COUNTERS and numeric(value)]
+
+
+class TensorBoard(Logger):
+    def __init__(self, dir=None, updates=True, flush_every=1):
+        self.dir = dir
+        self.updates = bool(updates)
+        self.flush_every = positive_int(flush_every, "flush_every")
+        self.writer = None
+        self.writer_class = None
+        self.turns = 0
+
+    def bind(self, ctx):
+        self.writer_class = summary_writer()
+
+    def directory(self, ctx):
+        if self.dir is not None:
+            return str(self.dir)
+        record = getattr(ctx, "record", None)
+        if record is None:
+            where = f"{self.path}: " if self.path else ""
+            raise BuildError(f"{where}the tensorboard logger needs a dir when the run writes no record")
+        return str(record.path("tensorboard"))
+
+    def open(self, ctx):
+        if self.writer is None:
+            if self.writer_class is None:
+                self.writer_class = summary_writer()
+            self.writer = self.writer_class(log_dir=self.directory(ctx))
+        return self.writer
+
+    def turn(self, ctx, line):
+        writer = self.open(ctx)
+        step = int(line["turn"])
+        for key in series_of(line):
+            if self.updates and per_update(key):
+                continue
+            writer.add_scalar(key, float(line[key]), step)
+        self.turns += 1
+        if self.turns % self.flush_every == 0:
+            writer.flush()
+
+    def update(self, ctx, line):
+        if not self.updates:
+            return
+        writer = self.open(ctx)
+        step = int(line["update"])
+        for key in series_of(line):
+            if per_update(key):
+                writer.add_scalar(key, float(line[key]), step)
+
+    def close(self, ctx):
+        if self.writer is None:
+            return
+        writer, self.writer = self.writer, None
+        writer.flush()
+        writer.close()
+
+
+class Echo(Logger):
+    def __init__(self, series=None, stream=None):
+        self.series = None if series is None else [series] if isinstance(series, str) else list(series)
+        self.stream = stream
+
+    def turn(self, ctx, line):
+        keys = series_of(line) if self.series is None else self.series
+        pairs = [f"{key} {line[key]:.4g}" for key in keys if numeric(line.get(key))]
+        stream = sys.stderr if self.stream is None else self.stream
+        stream.write("  ".join([f"turn {line['turn']}", *pairs]) + "\n")
+        flush = getattr(stream, "flush", None)
+        if callable(flush):
+            flush()
+
+
+register("tensorboard", "kalfa.std.loggers.TensorBoard", slot="logger",
+         doc="writes every history series per turn and the update losses to <record>/tensorboard for TensorBoard; "
+             "needs the tensorboard package (uv sync installs the dependencies of kalfa)")
+register("echo", "kalfa.std.loggers.Echo", slot="logger",
+         doc="prints the named series (or all) per turn to stderr; a logger to copy when writing your own")

kalfa/std/losses.py

@@ -0,0 +1,161 @@
+import torch
+from torch import nn
+
+from kalfa.core.aliases import register
+from kalfa.core.errors import RunError
+from kalfa.core.protocols import Objective
+from kalfa.core.wrap import loss_scalar
+from kalfa.std.common import classes_last_of
+
+
+def name_of(of):
+    return getattr(of, "name", "") or getattr(of, "__name__", "") or type(of).__name__
+
+
+def label_of(of, term):
+    name = name_of(of)
+    return name if term is None else f"{name}.{term}"
+
+
+def term_of(value, term, name):
+    if term is None:
+        return loss_scalar(value, name)
+    if isinstance(value, dict) and term in value:
+        return value[term]
+    have = f"its terms are {sorted(value)}" if isinstance(value, dict) else "it returns one value"
+    raise RunError(f"loss {name!r} has no term {term!r}; {have}")
+
+
+def items_of(items, kind, owner):
+    if isinstance(items, (dict, kind)):
+        items = [items]
+    out = []
+    for item in items or ():
+        if isinstance(item, kind):
+            out.append(item)
+        elif isinstance(item, dict):
+            out.append(kind(**item))
+        else:
+            raise TypeError(f"{owner} takes a list of {{of: ..., ...}} mappings, got {type(item).__name__}")
+    return out
+
+
+class Term:
+    def __init__(self, of, weight=1.0, term=None):
+        self.of = of
+        self.weight = float(weight)
+        self.term = term
+
+    def label(self):
+        return label_of(self.of, self.term)
+
+
+class WeightedSum(Objective):
+    def __init__(self, terms):
+        super().__init__()
+        self.terms = items_of(terms, Term, "weighted_sum")
+        if not self.terms:
+            raise ValueError('weighted_sum needs terms: a list of {of: "@losses.<name>", weight: <value>}')
+
+    def forward(self, outputs, batch, ctx):
+        out = {"loss": None}
+        total = None
+        for term in self.terms:
+            value = term_of(term.of(outputs, batch, ctx), term.term, name_of(term.of))
+            out[term.label()] = value
+            total = term.weight * value if total is None else total + term.weight * value
+        out["loss"] = total
+        return out
+
+
+class Constraint:
+    def __init__(self, of, epsilon, term=None, scale=1.0, damping=1.0, lmbda_init=0.0):
+        self.of = of
+        self.epsilon = float(epsilon)
+        self.term = term
+        self.scale = float(scale)
+        self.damping = float(damping)
+        self.lmbda_init = float(lmbda_init)
+
+    def label(self):
+        return label_of(self.of, self.term)
+
+
+class Mdmm(Objective):
+    def __init__(self, objective, constraints, term=None):
+        super().__init__()
+        self.objective = objective
+        self.constraints = items_of(constraints, Constraint, "mdmm")
+        self.term = term
+        if not self.constraints:
+            raise ValueError('mdmm needs constraints: a list of {of: "@losses.<name>", epsilon: <value>}')
+        self.lmbda = nn.Parameter(torch.tensor([entry.lmbda_init for entry in self.constraints],
+                                               dtype=torch.float32))
+
+    def forward(self, outputs, batch, ctx):
+        value = term_of(self.objective(outputs, batch, ctx), self.term, name_of(self.objective))
+        out = {"loss": None, "objective": value}
+        for position, constraint in enumerate(self.constraints):
+            held = term_of(constraint.of(outputs, batch, ctx), constraint.term, name_of(constraint.of))
+            infeasibility = constraint.epsilon - held
+            value = value + constraint.scale * (self.lmbda[position] * infeasibility
+                                                + constraint.damping * infeasibility ** 2 / 2)
+            label = constraint.label()
+            out[f"lambda/{label}"] = self.lmbda[position]
+            out[f"inf/{label}"] = infeasibility
+        out["loss"] = value
+        return out
+
+
+class CrossEntropy(nn.CrossEntropyLoss):
+    def __init__(self, weight=None, label_smoothing=0.0, classes_last=None):
+        super().__init__(weight=weight, label_smoothing=label_smoothing)
+        self.classes_last = classes_last
+
+    def forward(self, pred, target):
+        if classes_last_of(pred, target, self.classes_last, self, "cross_entropy") and pred.ndim >= 3:
+            pred = pred.movedim(-1, 1)
+        return super().forward(pred, target)
+
+
+def cross_entropy(weight=None, label_smoothing=0.0, classes_last=None):
+    if weight is not None:
+        weight = torch.as_tensor(weight, dtype=torch.float32)
+    return CrossEntropy(weight=weight, label_smoothing=float(label_smoothing), classes_last=classes_last)
+
+
+LOSSES = (
+    ("mse", "torch.nn.MSELoss", "mean squared error"),
+    ("l1", "torch.nn.L1Loss", "mean absolute error (the short name mae is the metric; write torch.nn.L1Loss or l1 "
+                              "for the loss)"),
+    ("huber", "torch.nn.HuberLoss", "huber loss, quadratic below delta and linear above"),
+    ("bce_logits", "torch.nn.BCEWithLogitsLoss", "binary cross entropy over logits"),
+    ("poisson_nll", "torch.nn.PoissonNLLLoss", "poisson negative log likelihood"),
+    ("cross_entropy", "kalfa.std.losses.cross_entropy", "cross entropy over logits (classes in dim 1 or last, as a "
+                                                        "language model gives them) with optional class weights "
+                                                        "(for example @data.class_weights) and label smoothing; "
+                                                        "classes_last: true or false states the layout when the "
+                                                        "tokens and the vocabulary have the same size"),
+    ("weighted_sum", "kalfa.std.losses.WeightedSum", "a weighted sum of other losses: terms: [{of: @losses.a, "
+                                                     "weight: 1.0, term: <key of a dict result>}]"),
+    ("mdmm", "kalfa.std.losses.Mdmm", "modified differential method of multipliers: an objective plus constraints "
+                                      "[{of: @losses.b, epsilon, scale, damping, lmbda_init, term}] with a "
+                                      "lambda parameter per constraint; give lmbda to an optimizer group with a "
+                                      "negative lr"),
+)
+
+OPTIMIZERS = (
+    ("adam", "torch.optim.Adam", "adam"),
+    ("adamw", "torch.optim.AdamW", "adam with decoupled weight decay"),
+    ("sgd", "torch.optim.SGD", "stochastic gradient descent with optional momentum"),
+)
+
+
+def register_names():
+    for name, target, doc in LOSSES:
+        register(name, target, slot="loss", doc=doc)
+    for name, target, doc in OPTIMIZERS:
+        register(name, target, slot="optimizer", doc=doc)
+
+
+register_names()