🚀 1. ネットワークを走らせる¶
最小のプログラムを走らせ、1 行ずつ説明する。2 つのノードで \(y = 2x + 1\) を計算するプログラムである。後半ではネットワークを分岐させ、合流させる。
flowchart LR
H1["ホスト"] -->|"x を書く"| S["scale（2 倍）"]
S --> F["shift（+1）"]
F -->|"y を読む"| H2["ホスト"]プログラム¶
examples/quickstart.py
"""The smallest program: two nodes in a network, running in the background.
The host writes an input, waits until the output arrives, and stops.
"""
from typing import override
import torch
import mujou
from mujou import Network, Node, Port, TensorSpec
VECTOR = TensorSpec((3,)) # every message here is a float32 vector of 3
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
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
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")
リポジトリのルートで次のように実行する。
次のような 1 行が表示される。
y は x = [1, 2, 3] に対する \(2x + 1\) である。t は shift が y を送った時刻で、単位は実行のタイマーの秒。値は実行ごとに変わる。
ノード¶
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
- ノードは
mujou.Nodeを継承したクラス。どのサブクラスとも同じく、__init__では最初にsuper().__init__()を呼び、各 hook でも最初にsuper()の同じ hook を呼ぶ __init__でノードの Port を属性として宣言する。Port(recv=VECTOR)は長さ 3 の float32 のベクトルを受け取り、Port(send=VECTOR)はそれを送る。慣習では前のノード側の Port をprev、次のノード側をnextと名付けるaとbは設定。引数から決めた数値で、状態ではないon_update(dt)はノードの 1 回分の処理。mujou がおよそ 10 ms ごとに繰り返し呼ぶ。recv()は相手が送った最新の値を待たずに返し、send()はすべての相手にコピーを送る
ネットワーク¶
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
ネットワークはノードを属性として持ち、その Port を connect でつなぎ、外から使う Port を自分の属性として公開する（self.prev、self.next）。connect は互いの送る spec が相手の受け取る spec と一致するかを確かめ、一致しなければ TypeError を送出する。
実行¶
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")
mujou.run(net)はネットワークを検査して全ノードを起動し、動き出したらRunを返す。各ノードは実行のワーカースレッドの上でon_updateを繰り返し、ホストのスレッドは先へ進むrun.writeはつながっていない Portnet.prevにxを書く。値は次に書くまで残るrecv()は待たないので、xがscaleを通り抜けるまでshiftは前の入力（最初はゼロ）から計算した値を送り続ける。run.wait_forは条件を満たした最初の出力を時刻と一緒に返す。実時間で 5 秒たつとTimeoutErrorを送出するwithブロックを抜けると実行が終了する。ノードの例外も遅くともそこでホストに届く
分岐と合流¶
送るだけの Port は相手をいくつでも取れる。複数の Port につなげば、送った値はそれぞれに届く（fan-out）。合流には MultiPort を使う。その recv() は Channel ごとに値を 1 つずつ、connect の順の dict[Channel, 値] で返し、どう束ねるかはノードが決める。
class Apply(Node):
"""Sends ``fn(x)``."""
def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> None:
super().__init__()
self.fn = fn
self.prev = Port(recv=VECTOR)
self.next = Port(send=VECTOR) # only sends: any number of peers
@override
def on_update(self, dt: float) -> None:
super().on_update(dt)
self.next.send(self.fn(self.prev.recv()))
class Sum(Node):
"""Sends the sum of the latest messages of all its peers."""
def __init__(self) -> None:
super().__init__()
self.prev = MultiPort(recv=VECTOR) # merges any number of peers
self.next = Port(send=VECTOR)
@override
def on_update(self, dt: float) -> None:
super().on_update(dt)
received = self.prev.recv() # dict[Channel, torch.Tensor], in connect order
self.next.send(torch.stack(list(received.values())).sum(dim=0))
class Diamond(mujou.Network):
"""``shift -> (double, square) -> sum``: y = 2 (x + 1) + (x + 1)^2."""
def __init__(self) -> None:
super().__init__()
self.shift = Apply(lambda x: x + 1)
self.double = Apply(lambda x: 2 * x)
self.square = Apply(lambda x: x**2)
self.sum = Sum()
self.shift.next.connect(self.double.prev) # one port, two peers
self.shift.next.connect(self.square.prev)
self.double.next.connect(self.sum.prev) # two peers, one MultiPort
self.square.next.connect(self.sum.prev)
self.prev, self.next = self.shift.prev, self.sum.next
Diamond は \(y = 2(x + 1) + (x + 1)^2\) を計算する。shift が両方の枝へ送り、sum が 2 つの枝の値を足す。受け取る Port の recv() は値を 1 つ返すので、受け取る相手は 1 つだけ。複数の相手から受け取るなら MultiPort を使う。uv run python examples/branching.py で実行できる。同じプログラムの後半は相手ごとに違う値を送る。説明は Port と Channel にある。