コンテンツにスキップ

🧩 2. ノードを書く

状態が微分方程式に従うノード、リーク積分器を書き、2 つを並べて走らせる。dt、状態、周期を扱う。どのノードも守る規則は Mujou の仕組みにある。

ノード

入力 \(u\) を時定数 \(\tau\) で追いかけるノード。状態 \(x\) は \(\frac{dx}{dt} = \frac{u - x}{\tau}\) に従う。

class Leaky(Node):
    """Follows its input ``u`` with the time constant ``tau`` (in seconds):

        dx/dt = (u - x) / tau

    ``on_update`` advances ``x`` by ``dt``, the time that passed since the last
    update, with the exact solution for an input held constant over the step:
    ``x`` moves the fraction ``1 - exp(-dt / tau)`` of the way to ``u``. That
    fraction stays below 1 for any ``dt``, so a long ``dt`` (a busy machine) never
    overshoots ``u``, and none of the elapsed time is lost.

    ``period`` is well below ``tau`` by default: a node downstream sees this
    output at most about one period late, a small delay next to ``tau``.
    """

    def __init__(self, tau: float, *, period: float = 0.001) -> None:
        super().__init__(period=period)
        self.tau = tau  # a setting, not state
        self.x = SCALAR.zeros()  # a tensor attribute: in the state by itself
        self.prev = Port(recv=SCALAR)
        self.next = Port(send=SCALAR)

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        u = self.prev.recv()
        self.x += (1 - math.exp(-dt / self.tau)) * (u - self.x)  # in place
        self.next.send(self.x)  # a copy: changing x later does not change it

このノードが書くのは次の 3 つ。

  1. __init__ で設定(tau)、状態(x)、Port(prev、next)
  2. on_update(dt) で入力を読み、状態を dt だけ進め、結果を送る
  3. それ以外は無い。on_update は mujou が繰り返し呼ぶので、ノードはループを持たない

dt に比例して更新する

dt はこのノードの前回の更新からの秒数で、実行のタイマーで測る。最初の dt は実行の開始から数える。数え始めるのは全ノードの on_start が終わった時点。

更新はいつも dt に比例させる。\(\frac{dx}{dt} = f(x)\) なら、最も簡単な更新は陽的 Euler 法の x += dt * f(x)。ノードは自分の速さで更新し、その速さは負荷で揺らぐ。更新を dt に比例させれば、どのノードも同じ時間の流れの上で同じダイナミクスに従い、マシンが混んでも dt が伸びるだけで済む。1 回あたりの刻みを固定すると、速く回るノードほど速く変化してしまう。

dt は一定ではなく、マシンが混むと長くなる。dt を切り捨てると、その分の時間が失われ連続時間の解からずれるので、切り捨てない。代わりに、起こりうる最も長い dt でも安定に更新する。方法は 2 つ。

  • 厳密解で進める: 厳密解がある系では、それで dt だけ進める。Leaky は \(x\) を \(u\) までの距離の 1 - exp(-dt / tau) 倍だけ動かす。\(u\) が一定の間は厳密で、係数はどんな dt でも 1 未満なので \(x\) は \(u\) を行き過ぎない
  • 刻みを分割する: 厳密解が無ければ、1 回の on_update の中で dt を安定な長さ以下の n 個の刻みに分け、for ループで n 回更新する(4. ノードの中で学習するの例)

不安定さがほかのノードとの結合から来るとき(読んだメッセージは dt の間そのまま)は、分割しても安定にならない。そのようなノードは、失う時間と引き換えに自分が使う dt を打ち切ってよい(mujou.pc のノードは max_dt で打ち切る)。

連続時間の更新に雑音を足すときは、標準偏差を dt ではなく \(\sqrt{dt}\) に比例させる。結果がノードの更新頻度に依存しなくなる。

状態

self.x はテンソルの属性なので、自動で状態に入る。run.state_dict() がこれを取得し、復元で書き戻す(6. 保存と再開)。状態は __init__ で作り、on_update ではその場で更新する(+=、copy_ など)。

属性に新しいテンソルを代入し直してもエラーにはならず、以後は新しいテンソルが状態になる。問題になるのは、ほかの属性や optimizer が最初のテンソルを参照している場合で、その参照は古いテンソルを指したまま残る。復元もノードが持つテンソルにその場で書き込む。

self.tau は引数から決めた数値で、状態ではなく設定。別の tau で作ったノードに復元しても、そのノードの設定のまま。実行中に変わる素の値(カウンタなど)は register_state で状態に加える(4. ノードの中で学習するの例)。

ノードは self.x を送った後でその場で書き換えているが、問題ない。send は値をコピーする。

周期

ノードは平均 period 秒に 1 回更新する。既定は 0.01 秒。on_update にかかった時間も周期に含み、遅れたノードが取り戻そうと続けて何回も更新することはない。Leaky の既定の周期は tau より十分短い 1 ms。後ろのノードは出力を最大で 1 周期ほど遅れて見るが、tau に比べれば小さい遅れで済む。

周期を決めても dt は period ちょうどにはならない。dt は実際に経った時間で、遅れると長く、その次は短くなる。周期を決めた場合も dt に比例して更新する。

周期を持つノード(センサーなど)から入力を受ける鎖では、後ろのノードは次の入力が来るまで同じ入力を読む。後ろのノードの周期は入力の周期以下にする。

period=0.0 は CPU のコアを使い切る

period=0.0 を渡すと、時刻が少しでも進めばすぐ次の on_update が呼ばれる。そのノードは CPU のコアを 1 つ使い切り、ほかのノードの計算時間を奪ってその dt を不規則にする。待ってはいけないノードにだけ渡す。

2 つ並べる

class LeakyChain(Network):
    """``u -> a -> b ->``: two leaky integrators with the time constant ``tau``."""

    def __init__(self, tau: float) -> None:
        super().__init__()
        self.a = Leaky(tau)
        self.b = Leaky(tau)
        self.a.next.connect(self.b.prev)
        self.prev, self.next = self.a.prev, self.b.next

leaky.py は鎖の入力を 0 から 1 に切り替え、各出力をそれが送られた時刻の厳密解と比べる。

if __name__ == "__main__":
    tau = 0.05
    net = LeakyChain(tau)
    with mujou.run(net) as run:
        # Step the input from 0 to 1 (write returns the time of the step), then
        # watch both integrators rise on the run's timer.
        start = run.write({net.prev: torch.tensor(1.0, dtype=torch.float64)})
        for _ in range(4):
            run.timer.sleep(tau / 2)
            a, t_a = run.read_with_time(net.a.next)
            b, t_b = run.read_with_time(net.b.next)
            # Each output against the exact one at the time it was sent. On a busy
            # machine a node may not have sent since the step: its time is then
            # earlier than start, and the difference is negative (None, before
            # its first send, counts as long before the step).
            exact_a, _ = step_response(-math.inf if t_a is None else t_a - start, tau)
            _, exact_b = step_response(-math.inf if t_b is None else t_b - start, tau)
            print(
                f"a = {a.item():.3f} (exact {exact_a:.3f}), "
                f"b = {b.item():.3f} (exact {exact_b:.3f})"
            )
        # Hold the input for 10 time constants: both reach the equilibrium x = u.
        run.timer.sleep(10 * tau)
        b = run.read(net.next)
    print(f"after {10 * tau:.1f} s more: b = {b.item():.4f}")
uv run python examples/leaky.py

run.write は切り替えた時刻を、run.read_with_time は各出力の時刻を返す。その差が実行のタイマーで測った切り替えからの時間で、ノードがどうスケジュールされても正しい。

次へ