kalfa 3.0.0.dev0__py3-none-any.whl

kalfa/std/samplers.py

@@ -0,0 +1,261 @@
+import itertools
+
+import torch
+from torch import nn
+
+from kalfa.core.aliases import register
+from kalfa.core.errors import RunError
+from kalfa.core.rng import Rng
+from kalfa.std.common import evaluating
+from kalfa.std.data.text import size_of
+from kalfa.std.objectives import NoiseTables
+
+
+STREAM = "samples"
+
+
+def module_of(model):
+    if isinstance(model, nn.Module):
+        return model
+    inner = getattr(model, "module", None)
+    if isinstance(inner, nn.Module):
+        return inner
+    if callable(model):
+        return model
+    raise RunError(f"a sampler needs a model to call, got {type(model).__name__}")
+
+
+def device_of(ctx, model):
+    device = getattr(ctx, "device", None)
+    if device is not None:
+        return Rng.place(device)
+    if isinstance(model, nn.Module):
+        for tensor in itertools.chain(model.parameters(), model.buffers()):
+            return tensor.device
+    return torch.device("cpu")
+
+
+def count_of(sampler, n):
+    count = int(sampler.n if n is None else n)
+    if count < 1:
+        raise RunError(f"{type(sampler).__name__} draws n samples; n must be positive, got {count}")
+    return count
+
+
+def rng_of(ctx):
+    return getattr(ctx, "rng", None)
+
+
+def normal(ctx, shape, device):
+    rng = rng_of(ctx)
+    if rng is None:
+        return torch.randn(shape, device=device)
+    return rng.normal(STREAM, shape, device=device)
+
+
+def first(result):
+    if isinstance(result, dict):
+        result = next(iter(result.values()))
+    while isinstance(result, (tuple, list)):
+        result = result[0]
+    return result
+
+
+def finished(result):
+    result = first(result)
+    if isinstance(result, torch.Tensor):
+        return result.detach().cpu()
+    return result
+
+
+class GanSampler:
+    def __init__(self, model, latent, n_classes=None, n=64):
+        self.model = model
+        self.latent = int(latent)
+        self.n_classes = None if n_classes is None else int(n_classes)
+        self.n = int(n)
+
+    def sample(self, ctx, n=None):
+        count = count_of(self, n)
+        model = module_of(self.model)
+        device = device_of(ctx, model)
+        with evaluating(model):
+            z = normal(ctx, (count, self.latent), device)
+            if self.n_classes:
+                labels = torch.arange(count, device=device) % self.n_classes
+                return finished(model(z, labels))
+            return finished(model(z))
+
+
+class LatentSampler:
+    def __init__(self, model, latent, n=64):
+        self.model = model
+        self.latent = int(latent)
+        self.n = int(n)
+
+    def sample(self, ctx, n=None):
+        count = count_of(self, n)
+        model = module_of(self.model)
+        device = device_of(ctx, model)
+        with evaluating(model):
+            return finished(model(normal(ctx, (count, self.latent), device)))
+
+
+class DdpmSampler:
+    def __init__(self, model, schedule, shape, n=64, steps=None):
+        self.model = model
+        self.schedule = schedule
+        self.shape = tuple(int(part) for part in shape)
+        self.n = int(n)
+        self.steps = None if steps is None else int(steps)
+        self.noise = None
+
+    def tables(self, device):
+        source = self.schedule
+        if not isinstance(source, NoiseTables):
+            if self.noise is None or self.noise.schedule is not source:
+                self.noise = NoiseTables(source)
+            source = self.noise
+        placed = source.on(device)
+        return int(placed.steps), placed.betas, placed.alphas, placed.cumulative
+
+    def timesteps(self, total):
+        if self.steps is None:
+            return list(range(total - 1, -1, -1))
+        if self.steps < 1:
+            raise RunError(f"ddpm_sampler steps must be positive, got {self.steps}")
+        count = min(self.steps, total)
+        if count == 1:
+            return [total - 1]
+        spaced = torch.linspace(total - 1, 0, count, dtype=torch.float64).round().long().tolist()
+        return list(dict.fromkeys(int(step) for step in spaced))
+
+    def sample(self, ctx, n=None):
+        count = count_of(self, n)
+        model = module_of(self.model)
+        device = device_of(ctx, model)
+        total, betas, alphas, cumulative = self.tables(device)
+        size = (count, *self.shape)
+        schedule = self.timesteps(total)
+        with evaluating(model):
+            x = normal(ctx, size, device)
+            for position, step in enumerate(schedule):
+                t = torch.full((count,), step, dtype=torch.long, device=device)
+                predicted = first(model(x, t)).to(x.dtype)
+                if self.steps is None:
+                    x = self.ancestral(ctx, x, predicted, step, betas, alphas, cumulative, size, device)
+                else:
+                    following = schedule[position + 1] if position + 1 < len(schedule) else None
+                    x = self.implicit(x, predicted, step, following, cumulative)
+        return x.clamp(-1.0, 1.0).detach().cpu()
+
+    @staticmethod
+    def ancestral(ctx, x, predicted, step, betas, alphas, cumulative, size, device):
+        coefficient = betas[step] / (1.0 - cumulative[step]).sqrt()
+        x = (x - coefficient * predicted) / alphas[step].sqrt()
+        if step > 0:
+            x = x + betas[step].sqrt() * normal(ctx, size, device)
+        return x
+
+    @staticmethod
+    def implicit(x, predicted, step, following, cumulative):
+        current = cumulative[step]
+        previous = cumulative[following] if following is not None else torch.ones_like(current)
+        start = ((x - (1.0 - current).sqrt() * predicted) / current.sqrt()).clamp(-1.0, 1.0)
+        return previous.sqrt() * start + (1.0 - previous).sqrt() * predicted
+
+
+class LmSampler:
+    def __init__(self, model, tokenizer, prompt="", max_new_tokens=100, temperature=1.0, top_k=None, n=1):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.prompt = str(prompt)
+        self.max_new_tokens = int(max_new_tokens)
+        self.temperature = float(temperature)
+        self.top_k = None if top_k is None else int(top_k)
+        self.n = int(n)
+
+    def tokenizer_of(self, ctx):
+        tokenizer = self.tokenizer
+        if tokenizer is None or isinstance(tokenizer, str):
+            tokenizer = getattr(getattr(ctx, "data", None), "tokenizer", None)
+        if tokenizer is None or isinstance(tokenizer, str) or not callable(getattr(tokenizer, "decode", None)):
+            raise RunError("lm_sampler needs a tokenizer with encode and decode; write tokenizer: \"@data.tokenizer\"")
+        return tokenizer
+
+    @staticmethod
+    def window_of(ctx):
+        limit = getattr(getattr(ctx, "data", None), "seq_len", None)
+        if isinstance(limit, int) and not isinstance(limit, bool) and limit > 0:
+            return limit
+        return None
+
+    @staticmethod
+    def logits_of(result, size=None, length=None):
+        if isinstance(result, dict) and "logits" in result:
+            result = result["logits"]
+        result = first(result)
+        if not isinstance(result, torch.Tensor) or result.ndim not in (2, 3):
+            shape = tuple(result.shape) if isinstance(result, torch.Tensor) else type(result).__name__
+            raise RunError(f"lm_sampler needs logits of shape (batch, tokens, vocabulary) from the model, got {shape}")
+        if result.ndim == 3:
+            result = result[:, -1, :] if LmSampler.classes_last(result.shape, size, length) else result[:, :, -1]
+        return result.float()
+
+    @staticmethod
+    def classes_last(shape, size, length):
+        tokens, classes = shape[1], shape[2]
+        if tokens == length:
+            return True
+        if classes == length:
+            return False
+        if size is None or classes == size:
+            return True
+        if tokens == size:
+            return False
+        raise RunError(f"lm_sampler cannot find the vocabulary axis of logits of shape {tuple(shape)}: the model "
+                       f"read {length} tokens and the tokenizer has {size}; give (batch, tokens, vocabulary) or "
+                       "(batch, vocabulary, tokens)")
+
+    def pick(self, logits, generator):
+        if self.temperature <= 0:
+            return logits.argmax(dim=-1, keepdim=True)
+        logits = logits / self.temperature
+        if self.top_k:
+            keep = min(self.top_k, logits.shape[-1])
+            least = torch.topk(logits, keep, dim=-1).values[:, -1:]
+            logits = logits.masked_fill(logits < least, float("-inf"))
+        return torch.multinomial(torch.softmax(logits, dim=-1), 1, generator=generator)
+
+    def sample(self, ctx, n=None):
+        count = count_of(self, n)
+        tokenizer = self.tokenizer_of(ctx)
+        size = size_of(tokenizer)
+        model = module_of(self.model)
+        device = device_of(ctx, model)
+        prompt = [int(token) for token in tokenizer.encode(self.prompt)] if self.prompt else []
+        skip = 0 if prompt else 1
+        ids = torch.tensor([prompt or [0]] * count, dtype=torch.long, device=device)
+        limit = self.window_of(ctx)
+        rng = rng_of(ctx)
+        generator = None if rng is None else rng.stream(STREAM, device)
+        with evaluating(model):
+            for _ in range(self.max_new_tokens):
+                window = ids[:, -limit:] if limit else ids
+                logits = self.logits_of(model(window), size, window.shape[1])
+                ids = torch.cat([ids, self.pick(logits, generator)], dim=1)
+        return [tokenizer.decode(row) for row in ids[:, skip:].cpu().tolist()]
+
+
+REGISTRATIONS = (
+    ("gan_sampler", "kalfa.std.samplers.GanSampler",
+     "n samples of model(z) or model(z, label) with z from the samples stream and labels cycling over n_classes"),
+    ("latent_sampler", "kalfa.std.samplers.LatentSampler", "n samples of a decoder fed standard normal latents"),
+    ("ddpm_sampler", "kalfa.std.samplers.DdpmSampler",
+     "reverse diffusion from noise of the given shape with the beta schedule; steps limits the number of steps"),
+    ("lm_sampler", "kalfa.std.samplers.LmSampler",
+     "autoregressive text from the prompt: max_new_tokens draws at temperature, top_k limits the choices"),
+)
+
+for name, target, doc in REGISTRATIONS:
+    register(name, target, slot="sampler", doc=doc)

kalfa/std/samples.py

@@ -0,0 +1,72 @@
+import math
+from pathlib import Path
+
+import numpy
+import torch
+from PIL import Image
+
+from kalfa.core.record import atomic
+
+
+def image_grid(images, columns=8, pad=2):
+    images = torch.nan_to_num(torch.as_tensor(images).detach().float().cpu())
+    if images.ndim == 3:
+        images = images.unsqueeze(1)
+    count, channels, height, width = images.shape
+    columns = max(1, min(columns, count))
+    rows = math.ceil(count / columns)
+    low, high = float(images.min()), float(images.max())
+    if high > low:
+        images = (images - low) / (high - low)
+    else:
+        images = images.clamp(0.0, 1.0)
+    canvas = torch.ones((channels, rows * (height + pad) + pad, columns * (width + pad) + pad))
+    for index in range(count):
+        row, column = divmod(index, columns)
+        top = pad + row * (height + pad)
+        left = pad + column * (width + pad)
+        canvas[:, top:top + height, left:left + width] = images[index]
+    if channels == 1:
+        canvas = canvas.expand(3, -1, -1)
+    return canvas.permute(1, 2, 0).numpy()
+
+
+def write_png(grid, path, least=256, most=8):
+    pixels = numpy.clip(numpy.rint(numpy.asarray(grid) * 255.0), 0, 255).astype(numpy.uint8)
+    height, width = pixels.shape[:2]
+    scale = max(1, min(most, math.ceil(least / max(height, width, 1))))
+    image = Image.fromarray(numpy.ascontiguousarray(pixels))
+    if scale > 1:
+        image = image.resize((width * scale, height * scale), Image.Resampling.NEAREST)
+    with atomic(path) as temporary:
+        image.save(temporary, format="PNG")
+    return Path(path)
+
+
+def is_text(samples):
+    return isinstance(samples, (list, tuple)) and bool(samples) and all(isinstance(item, str) for item in samples)
+
+
+def is_images(samples):
+    return isinstance(samples, torch.Tensor) and samples.ndim == 4 and samples.shape[1] in (1, 3)
+
+
+def save_samples(samples, directory, stem, figures=None):
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    if isinstance(samples, str):
+        samples = [samples]
+    if isinstance(samples, torch.Tensor):
+        samples = samples.detach().cpu()
+    path = directory / f"{stem}.pt"
+    with atomic(path) as temporary:
+        torch.save(samples, temporary)
+    written = [path]
+    if is_text(samples):
+        text = directory / f"{stem}.txt"
+        with atomic(text) as temporary:
+            temporary.write_text("\n\n".join(samples) + "\n", encoding="utf-8")
+        written.append(text)
+    elif is_images(samples):
+        written.append(write_png(image_grid(samples), directory / f"{stem}.png"))
+    return written

kalfa/std/schedules.py

@@ -0,0 +1,74 @@
+import math
+
+from kalfa.core.aliases import register
+
+
+def fraction(update, steps):
+    return min(max(float(update) / float(steps), 0.0), 1.0)
+
+
+class Warmup:
+    def __init__(self, steps):
+        self.steps = int(steps)
+
+    def __call__(self, update):
+        if self.steps <= 0:
+            return 1.0
+        return fraction(update, self.steps)
+
+
+class Cosine:
+    def __init__(self, steps, warmup=0, floor=0.0):
+        self.steps = int(steps)
+        self.warmup = int(warmup)
+        self.floor = float(floor)
+
+    def __call__(self, update):
+        update = float(update)
+        if self.warmup > 0 and update < self.warmup:
+            return fraction(update, self.warmup)
+        if self.steps <= self.warmup:
+            return self.floor
+        progress = fraction(update - self.warmup, self.steps - self.warmup)
+        return self.floor + (1.0 - self.floor) * 0.5 * (1.0 + math.cos(math.pi * progress))
+
+
+class LinearWarmup:
+    def __init__(self, start, end, steps):
+        self.start = float(start)
+        self.end = float(end)
+        self.steps = int(steps)
+
+    def __call__(self, update):
+        if self.steps <= 0:
+            return self.end
+        return self.start + (self.end - self.start) * fraction(update, self.steps)
+
+
+class StepDecay:
+    def __init__(self, every, gamma=0.1):
+        self.every = int(every)
+        self.gamma = float(gamma)
+
+    def __call__(self, update):
+        if self.every <= 0:
+            return 1.0
+        return self.gamma ** (int(update) // self.every)
+
+
+SCHEDULES = (
+    ("warmup", "kalfa.std.schedules.Warmup", "a factor rising linearly from 0 to 1 over steps updates, then 1"),
+    ("cosine", "kalfa.std.schedules.Cosine", "linear warmup over warmup updates, then cosine decay from 1 to "
+                                             "floor at steps updates"),
+    ("linear_warmup", "kalfa.std.schedules.LinearWarmup", "a value moving linearly from start to end over steps "
+                                                          "updates, then end; a loss weight or an lr factor"),
+    ("step_decay", "kalfa.std.schedules.StepDecay", "a factor multiplied by gamma every every updates"),
+)
+
+
+def register_names():
+    for name, target, doc in SCHEDULES:
+        register(name, target, slot="schedule", doc=doc)
+
+
+register_names()

kalfa/std/steps.py

@@ -0,0 +1,133 @@
+import torch
+
+from kalfa.core.aliases import register
+from kalfa.core.errors import RunError
+from kalfa.core.loop import TrainTurn, autocast, batch_size, detached
+from kalfa.core.protocols import Step
+from kalfa.core.wrap import loss_scalar
+
+
+class Supervised(Step):
+    def __init__(self, order=None, amp=False, grad_clip=None, accumulate=1):
+        self.order = None if order is None else list(order)
+        self.amp = bool(amp)
+        self.grad_clip = None if grad_clip is None else float(grad_clip)
+        self.accumulate = int(accumulate or 1)
+        self.cached = (None, {})
+        self.pending = {}
+        self.stepped = True
+
+    def arrangement(self, ctx):
+        order = list(self.order or ctx.optimizers)
+        unknown = [name for name in order if name not in ctx.optimizers]
+        if unknown:
+            raise RunError(f"{self.path}.order names optimizers {unknown} that are not defined; the optimizers are "
+                           f"{sorted(ctx.optimizers)}")
+        return order
+
+    def outputs_for(self, batch, ctx):
+        entry = ctx.predicts()
+        if entry is None:
+            return {}
+        return entry.call(batch)
+
+    def minimize(self, entry, scalar):
+        if self.accumulate <= 1:
+            self.stepped = True
+            return entry.minimize(scalar, clip=self.grad_clip)
+        count = self.pending.get(entry.name, 0) + 1
+        if count == 1:
+            entry.zero_grad()
+        entry.backward(scalar / self.accumulate)
+        if count < self.accumulate:
+            self.pending[entry.name] = count
+            return None
+        self.pending[entry.name] = 0
+        norm = entry.clip(self.grad_clip)
+        entry.step()
+        self.stepped = True
+        return norm
+
+    def update(self, entry, outputs, batch, ctx):
+        if entry.minimizes is None:
+            raise RunError(f"optimizer {entry.name!r} minimizes nothing; write minimizes: \"@losses.<name>\"")
+        value = entry.minimizes(outputs, batch, ctx)
+        scalar = loss_scalar(value, entry.minimizes.name)
+        if not isinstance(scalar, torch.Tensor) or not scalar.requires_grad:
+            raise RunError(f"the loss {entry.minimizes.name!r} of optimizer {entry.name!r} carries no gradient; its "
+                           f"models are frozen or the loss does not depend on their parameters")
+        ctx.last_norms[entry.name] = self.minimize(entry, scalar)
+        return value
+
+    def train_batch(self, batch, ctx):
+        result = {}
+        self.stepped = False
+        for name in self.arrangement(ctx):
+            outputs = self.outputs_for(batch, ctx)
+            self.cached = (ctx.batches, detached(outputs))
+            result[name] = self.update(ctx.optimizers[name], outputs, batch, ctx)
+        return result
+
+    def eval_batch(self, batch, ctx):
+        stamp, outputs = self.cached
+        if stamp is not None and stamp == ctx.batches:
+            return outputs
+        return self.outputs_for(batch, ctx)
+
+    def predict(self, batch, ctx):
+        return self.outputs_for(batch, ctx)
+
+
+class Alternating(Supervised):
+    def __init__(self, order=None, repeat=None, fresh_batch=False, amp=False, grad_clip=None, accumulate=1):
+        super().__init__(order, amp, grad_clip, accumulate)
+        self.repeat = dict(repeat or {})
+        self.fresh_batch = bool(fresh_batch)
+
+    def step_once(self, entry, batch, ctx):
+        with autocast(ctx, self.amp):
+            outputs = self.outputs_for(batch, ctx)
+            value = self.update(entry, outputs, batch, ctx)
+        self.cached = (ctx.batches, detached(outputs))
+        return value
+
+    def train_turn(self, ctx):
+        turn = TrainTurn(ctx, self)
+        order = self.arrangement(ctx)
+        repeat = {name: int(self.repeat.get(name, 1)) for name in order}
+        stepped_by = {name: 0 for name in order}
+        while turn.more():
+            shared = None if self.fresh_batch else turn.next_batch()
+            if shared is None and not self.fresh_batch:
+                break
+            losses = {}
+            last = None
+            self.stepped = False
+            for name in order:
+                entry = ctx.optimizers[name]
+                for _ in range(repeat[name]):
+                    batch = turn.next_batch() if self.fresh_batch else shared
+                    if batch is None:
+                        break
+                    value = self.step_once(entry, batch, ctx)
+                    turn.means.record(entry.minimizes.name, value, batch_size(batch))
+                    losses[name] = value
+                    stepped_by[name] += 1
+                    last = batch
+            if last is None:
+                break
+            if self.stepped:
+                turn.record_update(losses)
+            turn.observe(last)
+        idle = [name for name in order if stepped_by[name] == 0]
+        if idle and turn.taken:
+            ctx.logger.warning(f"turn {ctx.turn}: optimizers {idle} took no step; the train loader ran out of "
+                               f"batches before their place in order {order}")
+        return turn.close()
+
+
+register("supervised", "kalfa.std.steps.Supervised", slot="step",
+         doc="One forward pass of the predicts model per batch, then every optimizer in order minimizes its loss")
+register("alternating", "kalfa.std.steps.Alternating", slot="step",
+         doc="Per update round each optimizer in order takes repeat[name] steps, on the shared batch or on a fresh "
+             "batch per step")

kalfa/std/triggers.py

@@ -0,0 +1,135 @@
+import math
+import time
+
+from kalfa.core.aliases import register
+from kalfa.core.protocols import Trigger
+
+
+def latest(ctx, key):
+    value = (getattr(ctx, "latest", None) or {}).get(key)
+    if value is None:
+        return None
+    number = float(value)
+    return None if math.isnan(number) else number
+
+
+class AfterEpoch(Trigger):
+    def __init__(self, at):
+        self.at = int(at)
+
+    def fires(self, ctx):
+        return ctx.turn >= self.at
+
+
+class Every(Trigger):
+    def __init__(self, epochs):
+        self.epochs = int(epochs)
+
+    def fires(self, ctx):
+        return self.epochs > 0 and ctx.turn % self.epochs == 0
+
+
+class Plateau(Trigger):
+    persist = ("best", "wait")
+
+    def __init__(self, watch, patience, mode="min", min_delta=0.0):
+        if mode not in ("min", "max"):
+            raise ValueError(f"mode must be min or max, got {mode!r}")
+        self.watch = str(watch)
+        self.patience = int(patience)
+        self.mode = mode
+        self.min_delta = float(min_delta)
+        self.best = None
+        self.wait = 0
+
+    def improved(self, value):
+        if self.best is None:
+            return True
+        if self.mode == "max":
+            return value > self.best + self.min_delta
+        return value < self.best - self.min_delta
+
+    def reset(self):
+        self.best = None
+        self.wait = 0
+
+    def fires(self, ctx):
+        value = latest(ctx, self.watch)
+        if value is None:
+            return False
+        if self.improved(value):
+            self.best = value
+            self.wait = 0
+        else:
+            self.wait += 1
+        return self.wait >= self.patience
+
+    def describe(self):
+        return f"{self.watch} plateau {self.patience}"
+
+
+class Above(Trigger):
+    def __init__(self, watch, value):
+        self.watch = str(watch)
+        self.value = float(value)
+
+    def fires(self, ctx):
+        current = latest(ctx, self.watch)
+        return current is not None and current > self.value
+
+    def describe(self):
+        return f"{self.watch} > {self.value}"
+
+
+class Below(Trigger):
+    def __init__(self, watch, value):
+        self.watch = str(watch)
+        self.value = float(value)
+
+    def fires(self, ctx):
+        current = latest(ctx, self.watch)
+        return current is not None and current < self.value
+
+    def describe(self):
+        return f"{self.watch} < {self.value}"
+
+
+class TimeBudget(Trigger):
+    persist = ("started",)
+
+    def __init__(self, minutes):
+        self.minutes = float(minutes)
+        self.started = None
+
+    def reset(self):
+        self.started = None
+
+    def fires(self, ctx):
+        now = time.time()
+        if self.started is None:
+            self.started = now
+        return (now - self.started) >= self.minutes * 60.0
+
+    def describe(self):
+        return f"after {self.minutes} minutes"
+
+
+TRIGGERS = (
+    ("after_epoch", "kalfa.std.triggers.AfterEpoch", "fires from the turn whose count reaches at (absolute, so a "
+                                                     "resumed run counts on)"),
+    ("every", "kalfa.std.triggers.Every", "fires every epochs turns"),
+    ("plateau", "kalfa.std.triggers.Plateau", "fires after patience turns without the watched series improving "
+                                              "by min_delta; turns without a value are not counted"),
+    ("above", "kalfa.std.triggers.Above", "fires when the watched series is above value"),
+    ("below", "kalfa.std.triggers.Below", "fires when the watched series is below value"),
+    ("time_budget", "kalfa.std.triggers.TimeBudget", "fires once minutes have passed since the first turn that "
+                                                     "asked it"),
+)
+
+
+def register_names():
+    for name, target, doc in TRIGGERS:
+        register(name, target, slot="trigger", doc=doc)
+
+
+register_names()