03 — Training

Training is not the graph’s. There is no g.fit(...): a Trainer takes a graph and trains it, so the same graph can be trained three ways without being touched, and a node is never asked whether it is being trained.

Three levels, and none knows the one above exists:

	scale
the graph	one `forward`
the `Trainer`	a training run — an afternoon
N training runs	a Python `for`. Notebook 4

import tempfile

import torch

import somatize.torch  # registers the codec that lets a tensor cross an edge
from somatize import Graph, Node, Opaque, Recorder, Store
from somatize.record import Live, curve, forwards, nodes, progress, spent
from somatize.torch import Trainer, freeze, parameters

torch.manual_seed(0)

<torch._C.Generator at 0x7f791428c9d0>

A node that holds weights is still just a node

It returns an Opaque, which carries something that only exists in this process — a tensor with its autograd graph attached. That is what keeps a backward pass possible: what crosses as data would have lost it.

It is wrapped on the way out and arrives unwrapped on the way in: the node below is handed a tensor. The wrapper is about crossing, not about what a node holds.

class Body(Node):
    def __init__(self, width=32):
        self.net = torch.nn.Sequential(
            torch.nn.Linear(8, width), torch.nn.ReLU(), torch.nn.Linear(width, width)
        )

    def forward(self, x, ctx):
        return Opaque(self.net(x))

    def parameters(self):
        return list(self.net.parameters())


class Head(Node):
    def __init__(self, width=32):
        self.out = torch.nn.Linear(width, 1)

    def forward(self, x, ctx):
        # `x` is a tensor, not an `Opaque`: what the node above wrapped is
        # unwrapped on the way in. The wrapper is about **crossing**, not about
        # what a node holds.
        return Opaque(self.out(x))

    def parameters(self):
        return list(self.out.parameters())


def built(width=32, lr=1e-2):
    g = Graph.somatize(Body(width).named("body") >> Head(width).named("head"))
    return g, torch.optim.Adam(parameters(g), lr=lr)

Data, made up on purpose

A regression whose answer is a fixed linear map plus noise, so the loss has somewhere to go and nothing has to be downloaded.

truth = torch.randn(8, 1)


def batch(how_many=64):
    x = torch.randn(how_many, 8)
    return x, x @ truth + 0.1 * torch.randn(how_many, 1)

One `step` is the primitive, `fit` is sugar over it

Forward, loss, backward, and an update when the group closes. Whatever does not fit in an epoch loop is written as a while over step.

g, optimizer = built()
store = Store(tempfile.mkdtemp())
live = Live(title="training")
recorder = Recorder(store, run="a-run", summarising=["loss"])

trainer = Trainer(
    g,
    objective=torch.nn.functional.mse_loss,
    optimizer=optimizer,
    watching=[recorder, live],
)

for _ in range(300):
    trainer.step(batch())

live

Figure from cell 7

The stream carries both vocabularies at once: the engine’s — which node ran, how long — and this level’s, loss and updated. They are not one type and never were: a loss is the trainer’s arithmetic and the engine cannot see it. What makes them one thing is the record they land in.

facts_of_a_step = __import__("somatize.record", fromlist=["facts"]).facts
facts_of_a_step(store, run="a-run", forward=10)

[{'began_us': '2', 'fact': 'ran', 'node': 'body', 'took_us': '100'},
 {'began_us': '123', 'fact': 'ran', 'node': 'head', 'took_us': '32'},
 {'fact': 'finished', 'took_us': '175'},
 {'fact': 'updated'},
 {'fact': 'loss', 'value': '1.2943092584609985'}]

progress(store, run="a-run")

Figure from cell 10

Where the time went

The aggregated view. With two nodes it is not a mystery; with twenty and a device hop it is the first thing to look at when a step is slower than it should be.

spent(store, run="a-run")

Figure from cell 12

A group of steps, one update

every=N is for a batch that does not fit but whose gradient does: N steps, one update, and each loss divided by N so the group pulls exactly as one step over the N batches would. micro=N is the other half — one step cut into N pieces — and the two multiply rather than compete.

g, optimizer = built()
seen = []
grouped = Trainer(
    g,
    objective=torch.nn.functional.mse_loss,
    optimizer=optimizer,
    every=4,
    watching=seen.append,
)
for _ in range(8):
    grouped.step(batch())

said = [one["fact"] for one in seen]
print("losses :", said.count("loss"))
print("updates:", said.count("updated"), "— eight steps, two updates")

losses : 8
updates: 2 — eight steps, two updates

Freezing, and why a cache needs it

.cached() says a node’s output is worth keeping. .frozen() says the node is settled, so what was kept is still what it would produce. Declaring a cache over something that can still change is refused before the first node runs — not later, as a net that quietly stopped training.

frozen_g = Graph.somatize(Body().named("body").frozen().cached() >> Head().named("head"))
# Declaring it is the graph's half; making it true is torch's, and without the
# digest of its weights two checkpoints would share one name.
freeze(frozen_g)
print("fingerprints:", {k: v[:16] + "…" for k, v in frozen_g.fingerprints().items()})

kept = tempfile.mkdtemp()
said = lambda f: print(" ", f["fact"], f.get("node", ""))

# `Opaque` on the way in too: a bare tensor is refused at the edge, and the
# refusal names the two right answers rather than picking one.
same = Opaque(torch.randn(4, 8))
for which in ("first time", "same batch again"):
    print(which)
    frozen_g.forward(same, store=kept, watching=said)

print("a different batch")
frozen_g.forward(Opaque(torch.randn(4, 8)), store=kept, watching=said)

fingerprints: {}
first time
  ran body
  kept body
  ran head
  finished
same batch again
  recalled body
  ran head
  finished
a different batch
  ran body
  kept body
  ran head
  finished

tensor([[0.0953],
        [0.0706],
        [0.1410],
        [0.1752]], grad_fn=<AddmmBackward0>)

The same batch again recalled the body instead of running it; a different batch did not, and could not — a key is made out of content, so a hit means the very same input reached the very same settled code. That is what makes a settled prefix worth declaring: it runs once per batch and is read back on every epoch after the first.

What the run left behind

export gives the weights node by node, and a Store keeps them. It is what a federated round is built out of — and level 3 has no type there either: a round is a for, and fedavg is a function.

digest = store.keep("a-run/weights", trainer.export())
print("kept as", digest)
print("nodes  :", sorted(store.recall("a-run/weights")))

kept as sha256:f8761a381f8af4786ae1d96fe759cfd8b85c5f7bda151a84dc5e9ee103124298
nodes  : ['body', 'head']

What is not here yet

Whether those gradients are healthy — dying, exploding, saturated, a layer nobody is updating. That is a diagnosis: an opinion about the record, with arguable thresholds, and it is the next slice. The invariant that keeps it honest is already written down: a diagnosis has to be reproducible from the stored record, without training again.