💾 6. 保存と再開

実行の状態を取得して保存し、新しい実行でその続きから再開する。状態が何を持つかは Mujou の仕組みを参照。

手順

"""Save the state of a run, and restore it in a new run that goes on from there.

A node collects its state by itself: every attribute that holds a
``torch.Tensor`` or an ``nn.Module``, and the state of its generator.
``Run.state_dict`` gathers the states of all nodes, the latest messages and the
time into one tree, a plain ``dict`` of CPU tensors. Save it with ``torch.save``
(or pickle), and read it back with ``torch.load``.
"""

import io
from typing import override

import torch
from torch import nn

import mujou
import mujou.time
from mujou import Node, Port, TensorSpec

VECTOR1 = TensorSpec((1,))


class Integrator(Node):
    """Integrates its input over time (dx/dt = u) and sends ``readout(x)``."""

    def __init__(self) -> None:
        super().__init__()
        self.x = VECTOR1.zeros()  # a tensor: part of the state
        self.readout = nn.Linear(1, 1)  # a module: part of the state
        self.prev = Port(recv=VECTOR1)
        self.next = Port(send=VECTOR1)

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        self.x += dt * self.prev.recv()  # in place: other references to x stay valid
        with torch.no_grad():  # computed only to send: no graph
            self.next.send(self.readout(self.x))


class Model(mujou.Network):
    def __init__(self) -> None:
        super().__init__()
        self.integrator = Integrator()
        self.prev = self.integrator.prev
        self.next = self.integrator.next


one = torch.ones(1)

# Run for a while, then stop and take the state at the end.
net = Model()
with mujou.run(net, inputs={net.prev: one}) as run:
    run.timer.sleep(0.2)
tree = run.state_dict()  # after a normal end: the state at the stop
saved = tree["nodes"]["integrator"]["state"]  # this node's state_dict()
print("state of the node:", {k: type(v).__name__ for k, v in saved.items()})
file = io.BytesIO()  # or a path, such as "state.pt"
# Every tensor of the tree is on the CPU, so torch.load needs no GPU (a CUDA node's
# generator state still restores only into a CUDA node).
torch.save(tree, file)

# A new instance of the same network goes on from the saved state and time. A
# network runs once: run the one the tree came from again, and mujou.run raises.
restored = Model()
file.seek(0)
loaded = torch.load(file)  # the default weights_only=True reads the plain dict
with mujou.run(restored, inputs={restored.prev: one}, state=loaded) as run:
    # On a timer of its own that starts at the time of the tree: the global timer
    # and the other runs do not go back.
    assert run.timer is not mujou.time.global_timer
    print(f"restored at t = {float(run.timer.now()):.3f} s of its own timer")
    run.timer.sleep(0.1)
after = run.state_dict()["nodes"]["integrator"]["state"]

x0, x1 = saved["x"].item(), after["x"].item()
print(f"x = {x0:.3f} at the save, {x1:.3f} after 0.1 s more")
# The integral went on from the saved value: from 0, it would be about 0.1. (How
# far it gets in 0.2 s and 0.1 s depends on when the last update falls.)
assert x1 > x0 > 0
# The module was restored too: the new instance had other random weights.
assert torch.equal(after["readout"]["weight"], saved["readout"]["weight"])

1. ノードは状態を属性に持つ。self.x（テンソル）と self.readout（モジュール）は、ノードの乱数生成器の状態と一緒に自動で集まる
2. 状態は __init__ で作り、on_update ではその場で更新する。こうすれば、optimizer などのほかの参照も状態を指したままになる（2. ノードを書く）
3. ホストは run.state_dict() で木を取得する。実行中でも、正常に終了した後（ここでは with ブロックの後）でもよい
4. 木を torch.save で保存する
5. 再開するときは同じ構成の新しいインスタンスを作り、木を mujou.run(net, state=tree) に渡す。ネットワークは 1 回だけ走るので、木を取得したネットワークは使えない（ノードが守る規則）

状態の木

木は CPU 上のコピーからなる素の dict で、ノードのパスをキーにする。各ノードの状態、各 Channel の最新のメッセージとその時刻、各ノードの最後の更新の時刻、タイマーの時刻を持つ。正確な形は API リファレンスの Run.state_dict にある。

取得する間は実行が一時停止する。各ノードは更新の途中では止まらず、2 回の更新の間で止まるので、木の中の値はすべて同じ時刻のものになる。止まっていた時間は以後の dt に入らない。この一時停止の前後で各ノードの on_pause と on_resume が呼ばれ、その間の I/O を止められる。タイマーの一時停止（run.timer.pause()）ではどちらも呼ばれない。

取得するとき	得られるもの
ネットワークの実行中	その時刻の状態
正常に終了した後	on_stop の後にノードが残した値。時刻は shutdown が終了を求めた時点
例外で終了した後	取得できず、state_dict は NetworkError を送出する。残したい状態は実行中に取得し、長い実行では定期的に取得する

再開するときの注意

• 再開した実行は専用のタイマーで走る。 木の時刻と時間の速さから始まる新しい Timer（run.timer）で走るので、mujou.time.global_timer やほかの実行の時刻は戻らない。state= と一緒に timer= を渡すと ValueError になる。global_timer.pause() では再開した実行は止まらないので、その実行の run.timer を止める
• 同じネットワークを作る。 再開するネットワークは、ノードのパス・Port・接続が木を取得したネットワークと同じでなければならない。違えば mujou.run は、足りない項目と余分な項目を列挙して ValueError を送出する。属性名を変えるとパスが変わるので、変える前に保存した木では再開できない。どちらもまだ何も起動していないので、木を直せば同じインスタンスで再開できる
• ノードごとの状態をそろえる。 ノードの状態のキーや shape が違うと、ノードの起動中に既定の load_state_dict が ValueError を送出する。そのため mujou.run は NetworkError（__cause__ がその ValueError）を送出し、そのインスタンスはもう走らない。新しく作り直す
• 入力。 state= と inputs= を両方渡すと、inputs の値が木の値を置き換える

自分の値を状態に入れる

自動で集まるのは、テンソル、mujou.Stateful のオブジェクト（どの nn.Module も、pc.clock.LearnClock も、mujou.Stateful を継承した自作のクラスも）、self.generator の状態。数値・文字列、テンソルのリストや dict は集まらない。state_dict と load_state_dict という名前のメソッドを持つだけのオブジェクトも集まらないので、mujou.Stateful を継承する。実行中に変わる素の値は、__init__ で属性に代入して登録する。

self.elapsed = 0.0
self.register_state("elapsed")

値はそのまま状態に入り、後で属性に何を代入しても状態に残る。登録できるのは代入済みの属性だけ（そうでなければ AttributeError）。値は既定の torch.load が読める型（int・float・bool・str・None とそのリスト・タプル・dict、テンソル）にする。

deque などそれ以外のものは、state_dict と load_state_dict の両方を override する。super().state_dict() の戻り値を拡張し、super().load_state_dict には自分で足したキーを除いて渡す。

from collections import deque
from collections.abc import Mapping
from typing import Any, override

import torch

from mujou import Node


class Recent(Node):
    """Keeps the last 8 inputs in a deque."""

    def __init__(self) -> None:
        super().__init__()
        self.recent: deque[torch.Tensor] = deque(maxlen=8)

    @override
    def state_dict(self) -> dict[str, Any]:
        return {**super().state_dict(), "recent": list(self.recent)}

    @override
    def load_state_dict(self, state: Mapping[str, Any]) -> None:
        self.recent = deque((x.clone() for x in state["recent"]), maxlen=8)  # 木の値は写して持つ
        super().load_state_dict({k: v for k, v in state.items() if k != "recent"})

state_dict はノードが一時停止している間に呼ばれる。状態の取得がほかのノードから見えないよう、その中では Port が受信も送信も断る。戻り値は mujou がコピーする。load_state_dict は on_start の後、最初の on_update の前に 1 回呼ばれる。

ファイル

• torch.save と、既定の weights_only=True の torch.load は、テンソル・数値・文字列・dict・list・tuple からなる木を読み書きできる。mujou が木に入れるのはこれだけ。独自クラスのオブジェクトが入った木（register_state や override から）は、この方法では読めない
• pickle.load と torch.load(..., weights_only=False) は、ファイルの中のコードを実行できる。信頼できないファイルをこれらで読まない
• 既定の load_state_dict は各テンソルをノードのデバイスへコピーするので、CUDA のノードから保存した木も GPU の無いマシンで読める。ただし乱数生成器の状態は CPU と CUDA で違うので、木を復元できるのは保存したノードと同じ種類のデバイスのノードだけ

複数の実行を同じ時刻で

run.state_dict() は 1 つの実行の状態を 1 つの時刻で取得する。1 つのタイマーで走る複数の実行の状態を同じ時刻で取得するには、先にタイマーを止める。手順は時間。

次へ

• 各部分の詳細はトピックへ。まずは Port と Channel