🧠 Mujou の仕組み¶
mujou でプログラムを書くのに要る 9 つの考え方を、出会う順に図とコードで説明する。最後の節はノードが守る規則。手を動かして試すのはチュートリアルで。
このページで分かること
- ノードに 1 回分の更新だけを書き、mujou がそれを繰り返し呼ぶ仕組み
- ノード同士が互いを待たず、最新の値だけでやり取りする仕組み
- その上に実行・時間・状態・テスト・学習がどう組み立てられているか
前提は PyTorch の nn.Module とテンソルを使ったことがあること。
全体像¶
flowchart TB
H["ホスト(ネットワークを動かす Python のコード)"]
subgraph N["ネットワーク"]
direction LR
A["ノード A"] <-->|Channel| B["ノード B"] <-->|Channel| C["ノード C"]
end
T["実行のタイマー"]
H -->|"mujou.run(net) が返す Run で、入力を書き、出力を読み、状態を取得する"| N
T -.->|"時刻と dt"| N
図の部品と、それを説明する節の対応。
| 部品 | ひとことで | 節 |
|---|---|---|
| ノード | 1 回分の更新を書くクラス。繰り返し呼ぶのは mujou | 1. ノード |
| Port と Channel | ノードの接続口と、2 つの口をつなぐ線 | 2. Port と Channel |
| spec | Port が送受信するテンソルの shape と dtype | 3. メッセージの型 |
| ネットワーク | ノードを組み立てる入れ物。中のノードはパスで指す | 4. ネットワークとパス |
| 実行 | mujou.run が全ノードを動かし、ホストに Run を返す |
5. 実行 |
| 時間 | 実行ごとに 1 本のタイマー。更新は経過時間 dt に比例 |
6. 時間 |
| 状態 | 実行全体の状態を 1 つの dict で取得し、続きから再開 | 7. 状態 |
| テスト | 仮想のタイマーで時間を手で進め、決定的に走らせる | 8. テスト |
| 学習 | 更新則はノードの中。学習と推論のモードもエポックも無い | 9. 学習 |
PyTorch との対応¶
PyTorch を知っていれば、次の対応から入ると早い。
| PyTorch | mujou |
|---|---|
モデルの forward を自分のループで呼ぶ |
ノードの on_update(dt) を mujou が呼ぶ。ループは書かない |
forward の引数と戻り値 |
Port の recv() と send() |
| 親のモジュールの属性に子を代入する | Network の属性にノードを代入する |
model.train() と model.eval() |
無い。推論と学習を同時に行う |
loss.backward() がモデル全体に届く |
勾配はノードの中だけ。ノードの間はメッセージ |
model.state_dict() |
ノードごとの state_dict() と、実行全体の run.state_dict() |
1. ノード¶
要点
ノードは 1 回分の更新を on_update(dt) に書くクラス。ループは書かず、繰り返し呼ぶのは mujou。
ノードは計算の単位で、mujou.Node を継承したクラスとして書く。次は受け取った値を a * x + b にして送るノード。
class Affine(Node):
"""Sends ``a * x + b``, where ``x`` is the latest input."""
def __init__(self, a: float, b: float) -> None:
super().__init__()
self.a, self.b = a, b # settings, not state
self.prev = Port(recv=VECTOR) # receives from the previous node
self.next = Port(send=VECTOR) # sends to the next node
@override
def on_update(self, dt: float) -> None:
super().on_update(dt)
self.next.send(self.a * self.prev.recv() + self.b) # recv() never waits
__init__ で Port を作り、on_update(dt) で読み、計算し、送る。ノードの状態は自分の属性に持つ。
hook の呼ばれる順¶
on_update のほかの hook も mujou が呼ぶ。1 つの実行の中でノードが通る主な流れは次のとおり。
stateDiagram-v2
[*] --> on_start
on_start --> on_update : 復元するときは、ここで状態を読み込む
on_update --> on_update : 平均 period ごと
on_update --> on_pause : run.state_dict()
on_pause --> on_resume : 状態を複製した後
on_resume --> on_update
on_update --> on_stop : 正常な終了
on_update --> on_exception : どれかのノードの hook が例外
on_exception --> on_stop
on_stop --> [*]
| hook | 呼ばれるとき |
|---|---|
on_start() |
実行の始まりに 1 回(復元するときも)。I/O はここで開く |
on_update(dt) |
平均 period ごとに繰り返し。dt は前回の更新からの秒数 |
on_pause()・on_resume() |
run.state_dict() が状態を取得するための停止と、その再開 |
on_exception() |
どれかのノードの hook が例外を出したとき、全ノードで on_stop の前に 1 回 |
on_stop() |
実行が止まるときに 1 回(例外の後も)。I/O などの資源を解放する |
正確な順と細かい条件は API リファレンスの Node にある。
override した hook では最初に super() を呼ぶ
__init__ でも hook でも、先頭で super().__init__(...) や super().on_update(dt) を呼ぶ。上のコードもそう書いている。
内側は同期、間は非同期¶
ノードの境界で、計算の仕方が切り替わる。
| ノードの内側 | ノードの間 | |
|---|---|---|
| 実行 | 同期。ふつうの PyTorch のコード | 非同期。互いを待たない |
| やり取り | 変数と計算グラフ | Port のメッセージ(計算グラフから切り離したコピー) |
| 勾配 | 通常の PyTorch と同じく記録する | 届かない |
1 つのノードにどれだけの計算を入れるかは、ノードのクラスで決まる。1 つの層でも、複数の層でも、モデル全体でもよい。一緒に計算したい層(BP で一緒に学習させたい層も)は 1 つのノードのクラスにまとめる。
flowchart LR
subgraph N1["ノード 1(BP で一緒に学ぶ 2 層)"]
direction LR
L1["Linear"] --> R["ReLU"] --> L2["Linear"]
end
subgraph N2["ノード 2"]
L3["Linear"]
end
N1 -. "メッセージ(非同期)" .-> N2
実線はノードの内側で同期して進む計算、点線はノードをまたいで非同期に届くメッセージ。
ノードにループを書かないのはなぜ?
mujou は多くのノードを少ないワーカースレッドで回す。ノードが 1 回分の更新だけを書けば、いつ・どのスレッドで呼ぶかを mujou が決められる。状態を取得するために全ノードを止めることも、テストで 1 回ずつ進めることも、この形だからできる。
2. Port と Channel¶
要点
Port はノードの接続口、Channel は 2 つの Port をつなぐ線。読む側は待たず、相手が最後に送った値を読む。
connect が 2 つの Port を 1 本の Channel でつなぎ、その Channel を返す。
値は溜まらず、最新の値だけが残る¶
Channel はキューではない。送った値は前の値を上書きし、recv() はそのときの最新の値を返す。
sequenceDiagram
participant A as ノード A(送る)
participant C as Channel
participant B as ノード B(読む)
B->>C: recv()
C-->>B: 初期値(既定はゼロ、時刻は None)
A->>C: send(1)
A->>C: send(2)
B->>C: recv()
C-->>B: 2(1 は読まれずに上書きされた)
B->>C: recv()
C-->>B: 2(同じ値をもう一度)
A->>C: send(3)
- 読むときは待たない。相手がまだ送っていなくても、すぐに値が返る
- 相手が最初に送るまでは、相手の Port の初期値(既定はゼロ)を返し、その時刻は
None - 値には送った時刻が付く。
recv_with_time()で時刻と一緒に読めば、新しい値か同じ値の読み直しかがわかる
キューにしないのはなぜ?
キューに値を溜めると、送る側と読む側の速さの差で値が溜まっていき、遅い側に合わせてノード同士の歩調がそろってしまう。最新の値だけを読み書きすれば、各ノードは互いを待たずに自分の速さで動ける。
Port と MultiPort¶
flowchart LR
subgraph fan["Port: 送るだけなら何本でも"]
direction LR
S["送る Port"] --> R1["受け取る Port"]
S --> R2["受け取る Port"]
end
subgraph gather["MultiPort: 同じ種類の相手を束ねる"]
direction LR
X1["送る Port"] --> M["MultiPort<br/>recv() → dict[Channel, 値]"]
X2["送る Port"] --> M
end
Port |
MultiPort |
|
|---|---|---|
| 受け取る相手 | 1 つまで | いくつでも(同じ種類の相手) |
recv() の値 |
値 1 つ | dict[Channel, 値](connect の順) |
| 使いどころ | 1 対 1 のやり取り。送るだけなら 1 対多の配信 | 同じ種類の入力を束ねる。どう束ねるかはノードが決める |
役割の違う入力(予測と誤差など)は、MultiPort で束ねずに別々の Port で受ける。つなぎ方が循環してもよい。詳しくは Port と Channel。
3. メッセージの型(spec)¶
要点
Port は送るものと受け取るものの shape と dtype を spec で宣言する。spec が一致する Port どうしだけがつながる。
VECTOR = TensorSpec((3,)) # 要素 3 つの float32 のベクトル
self.prev = Port(recv=VECTOR) # 受け取るだけ
self.next = Port(send=VECTOR) # 送るだけ
送る側の spec が受け取る側の spec と一致するときだけ connect できる。値は torch.Tensor、Python のスカラー、またはその dict。
| spec | 運ぶもの | 使いどころ |
|---|---|---|
TensorSpec(shape, dtype) |
テンソル 1 つ | 活動や誤差のような 1 つの値 |
IntSpec(bits)、FloatSpec、BoolSpec、TimestampSpec |
int、float、bool、Timestamp 1 つ |
回数・フラグ・時刻。テンソルに包まない |
DictSpec |
TypedDict のテンソルとスカラーの組 |
入力とその目標のように、一緒に届けたい値 |
組になる値は 1 つのメッセージで送る
2 つの Port の値は別々に届く。読む側は一方の新しい値と、もう一方の古い値を同時に読むことがある。入力とその目標のように組になる値は DictSpec にまとめると、dict 全体が一度に届く。
send は計算グラフから切り離したコピーを送る。送った後で元のテンソルを変えても、送った値は変わらない。
4. ネットワークとパス¶
要点
Network はノードを組み立てる入れ物で、自分では何も実行しない。中のノードは、ルートからの属性名をドットでつないだパスで指す。
nn.Module に子のモジュールを持たせるのと同じように、ノードを属性に代入して Port をつなぐ。次は「2 倍する → 1 を足す」の 2 つのノードを並べたネットワーク。
class Pipeline(Network):
"""``x -> scale -> shift ->``: y = 2 x + 1."""
def __init__(self) -> None:
super().__init__()
self.scale = Affine(2.0, 0.0)
self.shift = Affine(1.0, 1.0)
self.scale.next.connect(self.shift.prev)
self.prev, self.next = self.scale.prev, self.shift.next # exposed as its own
最後の行で、最初のノードの入力と最後のノードの出力を自分の prev と next として公開している。こうすると、ネットワークもノードと同じように外からつなげる。
ネットワークを入れ子にすると、パスはルートからの属性名をつないだものになる。
flowchart TB
R["ルートのネットワーク"] -->|body| B["ネットワーク"]
R -->|head| H["ノード<br/>パス head"]
B -->|l1| L1["ノード<br/>パス body.l1"]
B -->|l2| L2["ノード<br/>パス body.l2"]
パスは状態のキーになる(7. 状態)。名前を付けて並べるだけなら、Group(**named) を使う。ネットワーク自身は何も実行せず、実行中も中のノードはそれぞれ独立に動く。
5. 実行¶
要点
mujou.run(root) が全ノードをバックグラウンドで動かし、ホストに窓口の Run を返す。ホストは Run で入力を書き、出力を読む。
ホストは利用者のプログラムのこと。上の Pipeline を走らせ、出力が届くまで待つ。
net = Pipeline()
x = torch.tensor([1.0, 2.0, 3.0])
# The nodes run on the run's worker threads until the end of the `with` block.
with mujou.run(net) as run:
run.write({net.prev: x})
# x takes a few updates to pass through the nodes: wait until y arrives, with
# the time it was sent (TimeoutError after 5 s of real time).
y, t = run.wait_for(net.next, lambda y: torch.equal(y, 2 * x + 1), timeout=5.0)
assert t is not None # None only for the zeros before the first send
print(f"y = {y} at t = {float(t):.3f} s")
sequenceDiagram
participant H as ホスト
participant R as Run
participant N as ノード(scale, shift)
H->>R: mujou.run(net)
Note over R,N: 木を検査し、全ノードを起動する
H->>R: write({net.prev: x})
loop 各ノードが平均 period ごとに
N->>N: on_update(dt)
end
H->>R: wait_for(net.next, 条件)
R-->>H: 条件を満たした値と、送った時刻
H->>R: with ブロックを抜ける
Note over R,N: 全ノードの on_stop を呼び、実行を終える
ホストが Run でできること。
| したいこと | 使うもの |
|---|---|
| 入力を書く | run.write |
| 最新の出力を読む | run.read・run.read_with_time |
| 条件を満たす出力を待つ | run.wait_for |
| 送られた値を順に受け取る | run.subscribe |
| 実行の時間で待つ | run.timer.sleep |
| 状態を取得する | run.state_dict() |
| 実行を終える | with ブロックを抜ける、または run.shutdown() |
ルートは 1 回だけ走る。もう一度走らせるときは新しいインスタンスを作る。終わった後は、ノードの属性をホストが読んでよい。詳しくはホストの入出力。
6. 時間¶
要点
1 つの実行の時間軸は 1 本のタイマー。ノードは前回の更新からの経過時間 dt に比例して状態を進める。
実行ごとに時間軸は 1 本で、それが実行のタイマーである。すべてのノードとホストがこれを使う。自分の Timer を渡さない実行は mujou.time.global_timer で走る。ホストは run.timer で時間を待ったり止めたりし、ノードは self.timer で読むだけ。
dt に比例して更新する¶
各ノードは平均 period ごとに 1 回更新する。ただし更新の間隔はノードごとに違い、揺らぐ。そのため on_update は回数ではなく、渡された dt に比例して状態を進める。
dt として受け取るたとえば入力 \(u\) に時定数 \(\tau\) で追従するノードは、次の微分方程式に従う。
on_update(dt) は、この式を dt だけ進める。
こう書くと、更新の速さが違い揺らぐノード同士でも同じダイナミクスに従う。マシンが混んでも dt が伸びるだけで、経った時間は失われない。このノード Leaky はノードを書くで作る。
2 種類の秒¶
| 秒の種類 | 使うもの |
|---|---|
| 実行のタイマーの秒 | dt、period、run.timer.sleep、値を送った時刻 |
| 実時間の秒 | timeout だけ |
ホストは実行のタイマーを一時停止したり、速さを変えたりできる。そのときも dt と period はタイマーに従い、timeout だけが実時間で測る。詳しくは時間。
7. 状態¶
要点
ノードの状態は属性から自動で集まる。run.state_dict() は実行全体の状態を 1 つの dict で取得し、mujou.run(..., state=tree) がその続きから走らせる。
ノードの状態になるもの¶
flowchart TD
A["ノードの属性"] --> Q1{"テンソル?"}
Q1 -- はい --> S["状態(自動で入る)"]
Q1 -- いいえ --> Q2{"state_dict と load_state_dict を持つ?<br/>(nn.Module など)"}
Q2 -- はい --> S
Q2 -- いいえ --> Q3{"実行中に変わる?"}
Q3 -- はい --> R["register_state で登録する"]
Q3 -- いいえ --> C["設定。状態に入らない"]
| 属性の例 | 状態になるか |
|---|---|
self.x = SCALAR.zeros() |
テンソルなので自動で入る |
self.model = nn.Linear(4, 2) |
mujou.Stateful なので自動で入る |
self.count = 0 |
self.register_state("count") で入る |
self.buffer = [](テンソルのリスト) |
self.register_state("buffer") で入る |
self.tau = 0.05(変わらない設定) |
入らない |
ノードの乱数生成器の状態も自動で入る。規則は API リファレンスの Node にある。
実行全体の状態を保存して再開する¶
run.state_dict() は実行を終えずに、その時点の実行全体の状態を取得する。各ノードの状態、各 Channel の最新のメッセージ、時刻を、パスをキーにした 1 つの素の dict にまとめたもの。
tree = run.state_dict()
# {
# "timer": {"time_ns": ..., "time_scale": ...},
# "nodes": {
# "body.l1": {"state": ..., "last_update_ns": ..., "ports": {...}},
# ...
# },
# }
flowchart LR
A["実行中の Run"] -->|"run.state_dict()"| T["状態の木<br/>(CPU の上の素の dict)"]
T -->|torch.save| F[("ファイル")]
F -->|torch.load| T2["状態の木"]
T2 -->|"mujou.run(新しいインスタンス, state=tree)"| B["続きから走る Run<br/>保存した時刻から始まる専用のタイマー"]
手順は保存と再開。
8. テスト¶
要点
mujou.testing.run は仮想のタイマーの上で、時間を手で進めて決定的に走らせる。
ふつうの実行では、各ノードがいつ更新するかは実行ごとに揺らぐ。テストでは時間と更新の順をテストが決める。
mujou.run |
mujou.testing.run |
|
|---|---|---|
| タイマー | 実行のタイマー(既定は global_timer) |
仮想のタイマー。run.step でだけ進む |
| 更新 | 各ノードが平均 period ごとに |
run.step(dt, *nodes) で 1 回ずつ。period は使わない |
dt |
実際に経った時間 | step で進めた時間 |
| 結果 | 実行ごとに揺らぐ | 毎回同じ |
run.step(dt, *nodes) は時間を dt だけ進めてから、渡したノードを渡した順に 1 回ずつ更新する。ノードを渡さなければ、全ノードを木の順に更新する。
sequenceDiagram
participant T as テスト
participant R as ManualRun
participant A as net.a
participant B as net.b
T->>R: run.step(h, net.a, net.b)
R->>R: 仮想のタイマーを h 進める
R->>A: on_update(h)
A-->>B: send(b はこの値を読む)
R->>B: on_update(h)
R-->>T: 戻る
次は 6 節の Leaky を厳密解と比べるテスト。
class TestLeaky:
def test_follows_the_exact_solution(self) -> None:
tau, h, steps = 0.1, 0.001, 100
node = Leaky(tau)
one = torch.tensor(1.0, dtype=torch.float64)
expected = 0.0
with mujou.testing.run(node, inputs={node.prev: one}) as run:
for _ in range(steps):
run.step(h) # 1 ms later, then every node once (period is ignored)
expected += (1 - math.exp(-h / tau)) * (1.0 - expected)
y, t = run.read_with_time(node.next)
assert t is not None and t.ns == steps * 1_000_000 # sent at the last step
torch.testing.assert_close(y, torch.tensor(expected, dtype=torch.float64))
手順はノードをテストする。
9. 学習¶
要点
更新則はノードの中に書く。ノードの間を流れるのはメッセージだけで、学習と推論のモードもエポックも無い。
学び続けるネットワーク¶
従来の機械学習は学習を終えたら重みを固定して推論する。mujou の走っているネットワークは入力を受け取り、出力を出し、同時に学ぶ。
flowchart LR
subgraph old["従来の機械学習"]
direction LR
T1["学習"] --> F["重みを固定"] --> I["推論"]
end
subgraph new["mujou"]
direction LR
IN["入力"] --> NW["走っているネットワーク"] --> OUT["出力"]
NW -->|"同時に学ぶ"| NW
end
更新則を選ぶ¶
mujou はノード・Port・スケジューラに更新則を組み込まない。どの則で学ぶかはノードが決める。
| 更新則 | 勾配を取る範囲 | 読むページ |
|---|---|---|
| ノードの中の BP | 1 つのノードの中。ふつうの loss.backward() |
ノードの中で学習する |
PC-ALM(mujou.pc) |
各ノードの中。隣とはメッセージでやり取り | PC-ALM で学習する |
| 自作の則 | ノードが決める | 更新則 |
flowchart LR
subgraph bp["ノードの中の BP"]
direction LR
X["入力"] --> M["1 つのノード<br/>forward → loss → backward → 更新"]
end
subgraph pc["PC-ALM"]
direction LR
P1["ノード"] -- 活動 --> P2["ノード"]
P2 -- 活動 --> P3["ノード"]
P3 -. エネルギーの勾配 .-> P2
P2 -. エネルギーの勾配 .-> P1
end
Port は計算グラフから切り離したコピーを送るので、ノードをまたぐ計算グラフは作れない。ノードの間を流れるのは活動・誤差・目標などのメッセージである。
ノードが守る規則¶
ほかのページに出てくるノードの注意は、ここにまとめてある。
| 規則 | 守らないと |
|---|---|
time.sleep でブロックしない |
ほかのノードの dt が伸びる |
| ノードは 1 回だけ走る | RuntimeError |
| 走っている間はホストからノードの属性を読まない | ノードの更新と競合する |
on_start では状態に依存する出力をしない |
復元する前の状態を送ってしまう |
on_stop で状態を変えない |
正常に終了した後に取得する状態が変わる |
| 実行中に変わる値は登録する | 復元したときに初めからになる |
| 計算グラフは 1 回の更新の中に閉じる | グラフが伸び続ける |
| 例外で終わった後は時刻が止まる | self.timer で測っても時間が経たない |
| 時間を動かせるのはホストだけ、ノードは Port だけで話す | RuntimeError |
スリープ¶
time.sleep でブロックしない。 ノードはワーカースレッドをほかのノードと共有し、ブロッキングする hook はその間ワーカーをふさぐ。ほかのノードの dt が伸びる。
一定間隔で更新するなら period を使う。デバイス・ソケット・lock を待つしかないノードには dedicated_thread=True を渡す。ホストは run.timer.sleep で待つ。こちらは実行の時間の速さと一時停止に従うが、time.sleep はどちらにも従わない。
再実行¶
ノードは 1 回だけ走る。 走ったノードやネットワークをもう一度走らせると RuntimeError になる。もう一度走らせるときは新しいインスタンスを作る。終わった後も属性は読める。
実行中の属性¶
走っている間はホストからノードの属性を読まない。 ノードの更新と競合する。Port を run.read で読むか、run.state_dict() を取得する。
on_start の出力¶
on_start では状態に依存する出力をしない。 復元するときは on_start の後に状態を読み込む。
on_stop と状態¶
on_stop で状態を変えない。 正常に終了した後に取得する状態は、on_stop の後にノードが残したもの。on_stop では I/O などの資源だけを解放する。
状態の登録¶
実行中に変わる値は登録する。 テンソルとモジュールは自動で集まる。素の値(カウンタ、最後に見た時刻)と、テンソルを入れたリストや dict(リプレイバッファなど)は register_state しないと、復元したときに初めからになる。
計算グラフの範囲¶
計算グラフは 1 回の更新の中に閉じる。 hook は通常の PyTorch と同じく勾配を記録する。パラメータの更新は torch.no_grad() の中で行い、パラメータから計算した値を属性で次の更新へ持ち越すときは detach する。そうしないとグラフが伸び続ける。
例外の後の時刻¶
例外で終わった後は時刻が止まる。 実行の時刻は終了の時点で止まるので、on_exception と on_stop の中では self.timer が進まない。ロボットを 1 秒かけて減速させるような処理は実時間で測る。
時間の操作と通信¶
時間を動かせるのはホストだけ、ノードは Port だけで話す。 hook の中で Timer.sleep・pause・resume・set_time_scale・VirtualTimer.advance や Run のメソッドを呼ぶと RuntimeError になる。1 つのノードが時間を動かすと、同じ実行のほかのすべてのノードの dt が変わってしまう。送った値は切り離されたコピーなので、backward() はほかのノードに届かない。