コンテンツにスキップ

🔥 4. ノードの中で学習する

ノードの内側の誤差逆伝播法（BP）で学ぶノードを書く。学ぶのは、実行中に目標が変わっていくストリームから。学習の段階も推論の段階も無く、ノードは新しいバッチが来るたびに予測し、学ぶ。

プログラム

教師のノードが、ホストの書く a を使って \(y = a x + b\) のバッチを 10 ms ごとに送る。ホストは実行の途中で a を変える。学習するノードの損失は下がり、目標が変わると跳ね上がり、また下がる。

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

学習するノード

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
  • モジュールは状態。 self.model はどの nn.Module とも同じく mujou.Stateful なので、テンソルの属性と同じく自動で状態に入る
  • seed。 seed=0 はノード専用の乱数生成器 self.generator の seed。nn.Linear の既定の初期化は全ノードが共有する生成器を使うので、重みは self.generator で初期化し直す
  • 素の値は register_state で状態に加える。 last と elapsed は実行中に変わるが、テンソルでもモジュールでもない。登録しておけばほかと一緒に保存・復元される
  • 新しいかどうか。 recv_with_time() はバッチを送られた時刻と一緒に返す。教師が最初に送るまで時刻は None、次に送るまで時刻は同じなので、学習するノードは各バッチから 1 回だけ学ぶ
  • 勾配は通常の PyTorch と同じ。 hook は通常の PyTorch と同じく勾配を記録する。loss.backward() が p.grad を埋め、パラメータは torch.no_grad() の中でその場で更新する。送るためだけの値（学習前の損失）もグラフは要らない
  • バッチの数ではなく時間で。 バッチごとの固定の刻みではなく、前に学んでからの時間（elapsed）だけ勾配流 \(\frac{dW}{dt} = -\text{rate}\,\frac{\partial L}{\partial W}\) をたどる。どれだけ学ぶかは、ノードが何回回ったかではなく時間で決まる。長い elapsed は MAX_STEP 以下の刻みに分け、陽的 Euler 法が発散しないようにする

BP はノードの中に閉じる。Port は計算グラフから切り離したコピーを送るので、勾配がノードをまたぐことはない。

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)

torch のスレッド数のようなプロセス全体の設定は、ホストが実行の前に決める。ノードが変えてよいのは自分のスレッドの設定だけで、hook を抜ける前に戻す（スレッドとデバイス）。

実行

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 は条件を満たすまで損失の Port を追う。終わった後は、ノードの属性をホストが読んでよい。learner.model が学習したモジュールである。

次へ