File size: 7,296 Bytes
ecc81b3 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """The shared shape of every model in the library.
A model here is nothing but: an optional input projection, a 1-D mixer type,
and an ``nd_method`` deciding how that mixer covers the lattice. One class
holds that recipe; ``LSTM``, ``GRU``, ``S4D``, and ``Mamba`` differ only in
``_mixer``. This is the design's central claim made literal — an N-D model is
a 1-D mixer plus a plan for sweeping it.
"""
from __future__ import annotations
import warnings
from collections.abc import Callable
from functools import partial
from typing import cast
import torch
import torch.nn as nn
from torch_dimensions.compose import axial_scan, resolve_nd_method
from torch_dimensions.lattice import AxisSpec, Lattice
from torch_dimensions.plan import ScanPlan
__all__ = ["LatticeModel"]
class LatticeModel(nn.Module):
"""Base for the model family. Subclasses set ``_mixer``.
With no lattice this is an ordinary sequence model — a lattice with no
spatial axes has an identity permutation, so the 1-D case is the N-D case
with nothing to fold.
"""
_mixer: type[nn.Module]
def __init__(
self,
d_model: int,
n_layers: int = 1,
lattice: Lattice | None = None,
*,
d_input: int | None = None,
nd_method: str | Callable[..., nn.Module] = axial_scan,
method: str | Callable[..., nn.Module] | None = None,
plan: ScanPlan | None = None,
bidirectional: bool | AxisSpec | list[AxisSpec] = False,
dropout: float = 0.0,
chunk: int | None = None,
mixer_kwargs: dict | None = None,
mixer: type[nn.Module] | None = None,
**method_kwargs,
) -> None:
super().__init__()
# `method` is the short spelling of `nd_method` — the method of
# multidimensionality. Both name the same thing; giving both is a
# contradiction waiting to happen and is refused.
if method is not None:
if nd_method is not axial_scan:
raise ValueError("pass either `method` or `nd_method`, not both")
nd_method = method
# No lattice means a single dynamic axis: an ordinary sequence.
self.lattice = lattice if lattice is not None else Lattice(shape=(), time=True)
if plan is None:
plan = ScanPlan.cyclic(self.lattice.axis_names, n_layers, bidirectional=bidirectional)
else:
if bidirectional is not False:
raise ValueError("pass either `plan` or `bidirectional`, not both")
# A plan *is* the layer schedule, so it fixes the depth, and it
# wins. But winning *silently* over a disagreeing n_layers would
# ship a model shallower (or deeper) than requested — the same
# silent-downgrade failure the schedule machinery exists to make
# loud. A warning rather than an error because generic builders
# legitimately fill n_layers unconditionally and add a plan only
# sometimes; n_layers=1 is the default and passes untouched.
if n_layers != 1 and n_layers != len(plan):
warnings.warn(
f"n_layers={n_layers} is ignored: the given plan has {len(plan)} steps "
"and a plan determines the depth",
UserWarning,
stacklevel=2,
)
# An input projection only when the data is not already d_model wide.
# Without it every caller writes the same nn.Linear, which is friction
# for no gain in purity.
self.in_proj = nn.Linear(d_input, d_model) if d_input is not None else nn.Identity()
# `mixer=` substitutes the 1-D operator without touching anything else,
# which is what makes a model debuggable: swapping in
# `td.testing.Recorder` answers "which axis did layer 3 sweep" on the
# real model rather than on a reconstruction of it. The class's own
# `_mixer` stays the default and the recorded config still names the
# class, so a substituted model is visibly not the stock one.
mixer_cls = mixer if mixer is not None else self._mixer
# A substituted mixer cannot be written into the recipe — a class is
# not serializable — so a checkpoint would silently rebuild with the
# stock one and hand back a *different model* that loads without
# complaint. `save()` refuses instead; see config._checkpoint_header.
self._substituted_mixer = None if mixer is None else mixer.__name__
self.nd = resolve_nd_method(nd_method)(
mixer=partial(mixer_cls, d_model, **(mixer_kwargs or {})),
plan=plan,
lattice=self.lattice,
d_model=d_model,
dropout=dropout,
chunk=chunk,
**method_kwargs,
)
# The construction recipe, recorded so a checkpoint can rebuild this
# model without the user re-specifying anything (td.save / td.load).
# Plain JSON-able types throughout — the validity mask becomes a
# nested list — so a config can also live in YAML unchanged. n_layers
# is recorded as the plan's true depth: the plan is the schedule, and
# a recipe that could disagree with itself would not be a recipe.
from torch_dimensions.config import lattice_to_dict, nd_method_name
self.config: dict = {
"d_model": d_model,
"n_layers": len(plan),
"lattice": None if lattice is None else lattice_to_dict(self.lattice),
"plan": plan.to_dict(),
"nd_method": nd_method_name(nd_method),
"d_input": d_input,
"dropout": dropout,
"chunk": chunk,
"mixer_kwargs": dict(mixer_kwargs) if mixer_kwargs else {},
**method_kwargs,
}
@property
def plan(self) -> ScanPlan:
return cast(ScanPlan, self.nd.plan)
def save(self, path) -> None:
"""Write this model — architecture and weights — to one checkpoint
file that :func:`torch_dimensions.load` can rebuild it from."""
from torch_dimensions.config import save
save(self, path)
def receptive_field(self) -> dict[str, dict[str, object]]:
"""How far along each axis this model can actually see.
``inf`` for mixers that span their axis in one layer (RNNs, SSMs,
attention); a finite number for convolutions, where it is a real
constraint worth checking before training rather than after. See
:func:`torch_dimensions.receptive_field`.
"""
from torch_dimensions.mixers.conv import axis_receptive_field
return axis_receptive_field(self)
def to_spec(self) -> dict:
"""A JSON-able description of this model's N-D architecture.
Derived without a forward pass, so it can be taken before any data
exists. See VIEWER.md.
"""
from torch_dimensions.spec import scan_model_spec
return scan_model_spec(self)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""``(B, [T,] *shape, d_input or d_model)`` in, ``d_model`` out.
With no lattice that is just ``(B, T, ...)``.
"""
return self.nd(self.in_proj(x))
|