Skip to content

🔥 4. Learn in a Node

This page writes a node that learns with backpropagation (BP) inside itself, from a stream whose target changes while it runs. There is no training phase and no inference phase: the node predicts and learns from each new batch as it comes.

The program

A teacher node sends a batch from \(y = a x + b\) every 10 ms, with the a the host writes; the host changes a while the network runs. The learner follows its loss: it falls, jumps when the target changes, and falls again.

uv run python examples/online_learning.py

t = 0.48 s: loss = 0.000090  (learned a = 2.0)
t = 0.49 s: loss = 3.405589  (a changed at 0.48 s)
t = 0.97 s: loss = 0.000095  (learned a = -1.0)
the learner's weight = -0.988, bias = 0.498

The learner

class Learner(Node):
    """Follows the gradient flow dW/dt = -rate * dL/dW, learning once from each
    new batch for the time since the last one, and sends its loss on each batch
    before learning from it."""

    MAX_STEP = 0.01  # seconds: a longer explicit Euler step could diverge

    def __init__(self, rate: float) -> None:
        # Updating more often than the teacher sends would only re-read its batch.
        super().__init__(period=0.005, seed=0)
        self.rate = rate  # a setting, not state
        self.model = nn.Linear(1, 1)  # has state_dict: in the state by itself
        # seed reproduces only self.generator: initialize the module from it.
        nn.init.normal_(self.model.weight, std=0.1, generator=self.generator)
        nn.init.zeros_(self.model.bias)
        self.last = -1  # the time (ns) of the batch learned from last
        self.elapsed = 0.0  # seconds since then
        self.register_state("last")  # plain values that change: register them
        self.register_state("elapsed")
        self.prev = Port(recv=BATCH_SPEC)
        self.loss = Port(send=SCALAR, initial=torch.tensor(math.inf))  # no batch yet

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        self.elapsed += dt
        batch, t = self.prev.recv_with_time()
        if t is None or t.ns == self.last:  # nothing sent yet, or nothing new
            return
        self.last = t.ns
        x, y = batch["x"], batch["y"]
        with torch.no_grad():  # the loss on an unseen batch needs no graph
            self.loss.send(nn.functional.mse_loss(self.model(x), y))
        # How far to learn is set by the time, not by the number of batches: move
        # by the elapsed time, in n equal steps no longer than MAX_STEP.
        n = max(1, math.ceil(self.elapsed / self.MAX_STEP))
        for _ in range(n):
            loss = nn.functional.mse_loss(self.model(x), y)  # records the graph
            loss.backward()
            with torch.no_grad():  # the step itself is not part of any graph
                for p in self.model.parameters():
                    assert p.grad is not None
                    p -= self.elapsed / n * self.rate * p.grad
                    p.grad = None
        self.elapsed = 0.0
  • The module is state. self.model is a mujou.Stateful, as every nn.Module is, so it is in the node's state by itself, as a tensor attribute is
  • The seed. seed=0 seeds self.generator, the node's own random generator. The default initialization of nn.Linear draws from the global generator that every node shares, so initialize the weights again from self.generator
  • Plain values join the state with register_state. last and elapsed change while the node runs, but they are neither tensors nor modules: registered, they are saved and restored with the rest
  • New or not. recv_with_time() returns the batch with the time it was sent. Its time is None until the teacher first sends, and it stays the same until the teacher sends again, so the learner learns once from each batch
  • Gradients as in plain PyTorch. The hooks record gradients as plain PyTorch does: loss.backward() fills p.grad, and the parameters are updated in place inside torch.no_grad(). Values computed only to send (the loss before learning) need no graph either
  • By the time, not by the batch. The learner follows the gradient flow \(\frac{dW}{dt} = -\text{rate}\,\frac{\partial L}{\partial W}\) over the time since it last learned (elapsed), not a fixed step per batch: how far it learns depends on time, not on how often it runs. A long elapsed is split into steps no longer than MAX_STEP, so that explicit Euler does not diverge

BP stays inside the node. A port sends a copy, detached from the graph, so no gradient crosses nodes.

Settings of torch

# The nodes compute on the run's worker threads, and torch adds threads of its
# own for each: keep those few (a setting of the whole process, made by the host).
torch.set_num_threads(1)

The host sets the settings of the whole process, such as the number of torch's threads, before the run. A node changes only settings of its own thread, and sets them back before its hook returns (see Threads and Devices).

The run

teacher, learner = Teacher(), Learner(rate=5.0)
teacher.next.connect(learner.prev)
net = mujou.Group(teacher=teacher, learner=learner)
with mujou.run(net, inputs={teacher.a: torch.tensor(A_BEFORE)}) as run:
    # How fast the learner learns depends on time, not on how often it runs:
    # follow its loss, giving each phase up to 5 s of real time.
    loss, t = run.wait_for(learner.loss, learned, timeout=5.0)
    assert t is not None  # the loss on a batch, not the initial inf
    print(f"t = {float(t):.2f} s: loss = {loss.item():.6f}  (learned a = {A_BEFORE})")
    # The world changes: the same learner keeps learning, with no new phase.
    changed_at = run.write({teacher.a: torch.tensor(A_AFTER)})
    loss, t = run.wait_for(learner.loss, lambda loss: loss.item() > 1e-2, timeout=5.0)
    assert t is not None
    print(
        f"t = {float(t):.2f} s: loss = {loss.item():.6f}  (a changed at {float(changed_at):.2f} s)"
    )
    loss, t = run.wait_for(learner.loss, learned, timeout=5.0)
    assert t is not None
    print(f"t = {float(t):.2f} s: loss = {loss.item():.6f}  (learned a = {A_AFTER})")

# After the run, the nodes' attributes are the host's to read.
weight, bias = learner.model.weight.item(), learner.model.bias.item()
print(f"the learner's weight = {weight:.3f}, bias = {bias:.3f}")

run.wait_for follows the loss port until the condition holds. After the run, the nodes' attributes are the host's to read: learner.model is the module the learner learned.

Next