jens/pyRBM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[StackRnn] - add unrolled (own-weights) mode; each position gets its own RBM

Browse Source 
Previously a single shared-weight Entity processed every time step.  Now:
- Shared mode  (1 layer via make_layer):  original behaviour unchanged
- Unrolled mode (N layers via make_unrolled): layers[t] owns W_t, b_v_t, b_h_t

New API:
  make_unrolled(T, sensory_size, h_size, ...)  → list[Layer]
  next_entity()   → Entity for the upcoming step() call
  current_entity() → Entity from the most recent step() call
  is_shared        → bool

moby_rnn.ipynb: switch Build model cell to make_unrolled(T=100); update
  predict_next() to use rnn.next_entity() for position-correct Gibbs sampling
README_moby_rnn.md: redraw temporal-unrolling ASCII art showing per-position
  weights W_t; update parameter count and checkpoint file listing

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jens
2026-05-31 11:52:01 +02:00
co-authored by Claude Sonnet 4.6
parent 0b302a9c12
commit 3698c85ed4
4 changed files with 229 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/rbm/stack_rnn.py
+144 -64
View File
@@ -16,33 +16,65 @@ _CD_FUNC = {
class StackRnn(Stack):
"""Recurrent RBM stack (RTRBM-style, concatenation variant).
"""Recurrent RBM stack.
At each time step t and layer i:
visible_i[t] = concat(h_i[t-1], input_i[t])
h_i[t] = entity_i.forward(visible_i[t])
Two modes selected by the number of appended layers:
where input_0[t] = x_t (sensory input) and input_{i+1}[t] = h_i[t].
Weights are shared across time — the same Entity processes every time step.
Shared weights (1 layer):
The same Entity processes every time step.
Sequences may have any length.
Unrolled / own weights (N layers):
layers[t % N] processes time step t — each position has its own W.
Training requires sequences of exactly length N.
"""
def __init__(self, name: str, work_dir: str = '.'):
Stack.__init__(self, StackType.Rnn, name, work_dir)
self._h: list[Mat] | None = None
self._h: Mat | None = None # context / hidden state: (batch, h_size)
self._t: int = 0 # current time-step counter
# ── Convenience factory ───────────────────────────────────────────────────
# ── Convenience factories ─────────────────────────────────────────────
@staticmethod
def make_layer(name: str, sensory_size: int, h_size: int,
entity_params: EntityParams, training_params: TrainingParams) -> Layer:
"""Create a Layer suitable for StackRnn.
The entity's visible size = sensory_size + h_size.
sensory_size is the raw input width; h_size is the recurrent state size.
"""
"""Single layer for shared-weights mode."""
return Layer(name, (1, sensory_size, h_size, h_size), entity_params, training_params)
# ── Derived sizes ─────────────────────────────────────────────────────────
@staticmethod
def make_unrolled(time_steps: int, sensory_size: int, h_size: int,
entity_params: EntityParams, training_params: TrainingParams) -> list[Layer]:
"""N layers for unrolled (own-weights) mode — one per time step.
Usage::
for layer in StackRnn.make_unrolled(T, ...):
rnn.append(layer)
"""
return [
Layer(f"t{t}", (1, sensory_size, h_size, h_size), entity_params, training_params)
for t in range(time_steps)
]
# ── Mode ─────────────────────────────────────────────────────────────
@property
def is_shared(self) -> bool:
return self.num_layers() == 1
def _entity_at(self, t: int) -> Entity:
return self.from_index(t % self.num_layers()).entity
def next_entity(self) -> Entity:
"""Entity that will be used by the next step() call."""
return self._entity_at(self._t)
def current_entity(self) -> Entity:
"""Entity used by the most recent step() call."""
return self._entity_at(max(0, self._t - 1))
# ── Derived sizes ─────────────────────────────────────────────────────
def h_size(self, layer_idx: int = 0) -> int:
return self.from_index(layer_idx).entity.shape[1]
@@ -51,55 +83,55 @@ class StackRnn(Stack):
e = self.from_index(layer_idx).entity
return e.shape[0] - e.shape[1]
# ── Hidden-state management ───────────────────────────────────────────────
# ── State management ──────────────────────────────────────────────────
def reset(self, batch_size: int = 1):
"""Zero all hidden states (call before processing a new sequence)."""
self._h = [np.zeros((batch_size, self.h_size(i))) for i in range(self.num_layers())]
"""Zero the context vector and reset the time-step counter."""
self._h = np.zeros((batch_size, self.h_size()))
self._t = 0
# ── Inference ─────────────────────────────────────────────────────────────
# ── Inference ─────────────────────────────────────────────────────────
def step(self, x: Mat) -> Mat:
"""One time step forward through all layers.
"""One time step forward.
In shared mode uses the single entity.
In unrolled mode uses layers[_t % N] and advances _t.
x: (batch_size, sensory_size) or (sensory_size,)
Returns the top-layer hidden state h.
Returns the new context vector h_t.
"""
if x.ndim == 1:
x = x[None, :]
batch_size = x.shape[0]
if self._h is None or self._h[0].shape[0] != batch_size:
if self._h is None or self._h.shape[0] != batch_size:
self.reset(batch_size)
x_in = _to_gpu(x)
for i, layer in enumerate(self.layers):
visible = np.concatenate([self._h[i], x_in], axis=1)
h_new = layer.entity.forward(visible)
self._h[i] = h_new
x_in = h_new
entity = self._entity_at(self._t)
visible = np.concatenate([self._h, _to_gpu(x)], axis=1)
self._h = entity.forward(visible)
self._t += 1
return self._h
return self._h[-1]
def reconstruct(self, h: Mat, layer_idx: int = -1) -> Mat:
def reconstruct(self, h: Mat) -> Mat:
"""Decode h → visible, returning only the sensory portion.
The visible layer is [context | sensory]; this method strips context,
returning only the sensory reconstruction.
h: (batch_size, h_size)
Uses the entity from the most recent step() call.
"""
if layer_idx < 0:
layer_idx = self.num_layers() + layer_idx
entity = self.from_index(layer_idx).entity
entity = self.current_entity()
visible = entity.reconstruct(h)
return visible[:, entity.shape[1]:]
# ── Training ──────────────────────────────────────────────────────────────
# ── Training ──────────────────────────────────────────────────────────
def train(self, sequences: Mat, status: Status = None):
"""Greedy layer-wise CD training over sequences.
"""Train on sequences.
sequences: (T, sensory_size) — single sequence
(num_seq, T, sensory_size) — batch of sequences
Shared mode (1 layer): T may be any value.
Unrolled mode (N layers): T must equal N.
"""
if status is None:
status = Status()
@@ -107,24 +139,31 @@ class StackRnn(Stack):
seqs = sequences if sequences.ndim == 3 else sequences[None, :]
num_seq, T, _ = seqs.shape
for layer_idx, layer in enumerate(self.layers):
entity = layer.entity
print(f"Train layer {layer_idx} ({entity.name}) for {entity.training_params.num_epochs} epochs")
if self.is_shared:
entity = self.from_index(0).entity
print(f"Train shared ({entity.name}) "
f"for {entity.training_params.num_epochs} epochs")
if entity.enable_training and entity.training_params is not None:
self._train_layer(entity, seqs, num_seq, T, status)
self._train_shared(entity, seqs, num_seq, T, status)
else:
assert T == self.num_layers(), (
f"Unrolled mode: sequence length T={T} "
f"must equal num_layers={self.num_layers()}"
)
params = self.from_index(0).entity.training_params
print(f"Train unrolled ({self.num_layers()} layers) "
f"for {params.num_epochs} epochs")
self._train_unrolled(seqs, num_seq, T, status)
if layer_idx < self.num_layers() - 1:
seqs = self._pass_through(entity, seqs, num_seq, T)
def _train_layer(self, entity: Entity, seqs: Mat,
num_seq: int, T: int, status: Status):
def _train_shared(self, entity: Entity, seqs: Mat,
num_seq: int, T: int, status: Status):
"""One entity, reused at every time step."""
cd_func = _CD_FUNC[entity.type]
params = entity.training_params
h_sz = entity.shape[1]
params = entity.training_params
h_sz = entity.shape[1]
d_progress = 100.0 / params.num_epochs
progress = 0.0
d_progress = 100.0 / params.num_epochs
progress = 0.0
keep_running = True
entity.grad_zero()
@@ -138,7 +177,7 @@ class StackRnn(Stack):
err_total = 0.0
for t in range(T):
x_t = _to_gpu(seqs[:, t, :])
x_t = _to_gpu(seqs[:, t, :])
visible = np.concatenate([h, x_t], axis=1)
dwhv, dbv, dbh = cd_func(entity, visible)
@@ -157,11 +196,10 @@ class StackRnn(Stack):
keep_running = False
break
# Final report using one clean forward pass
h = np.zeros((num_seq, h_sz))
err_total = 0.0
for t in range(T):
x_t = _to_gpu(seqs[:, t, :])
x_t = _to_gpu(seqs[:, t, :])
visible = np.concatenate([h, x_t], axis=1)
h = entity.forward(visible)
err_total += rms_error_accu(visible - entity.reconstruct(h))
@@ -170,15 +208,57 @@ class StackRnn(Stack):
"err_rms_total": {"value": err_total / T, "unit": ""},
})
def _pass_through(self, entity: Entity, seqs: Mat,
num_seq: int, T: int) -> _np_cpu.ndarray:
"""Run sequences through entity; return hidden outputs on CPU."""
h_sz = entity.shape[1]
outputs = _np_cpu.zeros((num_seq, T, h_sz), dtype=_np_cpu.float64)
def _train_unrolled(self, seqs: Mat, num_seq: int, T: int, status: Status):
"""N entities, one per time step — each has its own W, b_v, b_h."""
params = self.from_index(0).entity.training_params
h_sz = self.from_index(0).entity.shape[1]
d_progress = 100.0 / params.num_epochs
progress = 0.0
keep_running = True
for layer in self.layers:
layer.entity.grad_zero()
status.on_change(self.from_index(0).entity)
for epoch in range(params.num_epochs):
if not keep_running:
break
h = np.zeros((num_seq, h_sz))
err_total = 0.0
for t, layer in enumerate(self.layers):
entity = layer.entity
cd_func = _CD_FUNC[entity.type]
x_t = _to_gpu(seqs[:, t, :])
visible = np.concatenate([h, x_t], axis=1)
dwhv, dbv, dbh = cd_func(entity, visible)
grad = entity.grad_compute(dbv, dbh, dwhv)
entity.state_adjust(grad, 1.0 / num_seq)
h = entity.forward(visible)
err_total += rms_error_accu(visible - entity.reconstruct(h))
progress += d_progress
if status.want_report(round(progress)):
if not status.on_change(self.from_index(0).entity, {
"progress": {"value": round(progress), "unit": "%"},
"err_rms": {"value": err_total / T, "unit": ""},
}):
keep_running = False
break
h = np.zeros((num_seq, h_sz))
for t in range(T):
x_t = _to_gpu(seqs[:, t, :])
err_total = 0.0
for t, layer in enumerate(self.layers):
entity = layer.entity
x_t = _to_gpu(seqs[:, t, :])
visible = np.concatenate([h, x_t], axis=1)
h = entity.forward(visible)
outputs[:, t, :] = convert(h)
return outputs
err_total += rms_error_accu(visible - entity.reconstruct(h))
status.on_change(self.from_index(0).entity, {
"progress": {"value": 100, "unit": "%"},
"err_rms_total": {"value": err_total / T, "unit": ""},
})
src/tests/test_rnn.py
+51 -64
View File
@@ -1,8 +1,4 @@
"""Test for StackRnn: recurrent RBM with concatenated [h_{t-1} | x_t] visible layer.
Uses a repeating binary pattern as a minimal synthetic sequence so the model
has something learnable to compress and predict.
"""
"""Tests for StackRnn — shared-weights and unrolled (own-weights) modes."""
import numpy as _np_cpu
from rbm.stack_rnn import StackRnn
from rbm.matrix import np, convert
@@ -10,94 +6,85 @@ from rbm.entity import EntityParams, TrainingParams
from rbm.status import Status
SENSORY_SIZE = 8
H_SIZE = 4
T = 16 # sequence length
NUM_SEQ = 10 # sequences in the batch
WORK_DIR = "../../results"
H_SIZE = 4
T = 16 # sequence length
NUM_SEQ = 10 # sequences in the batch
WORK_DIR = "../../results"
_PARAMS = TrainingParams(learning_rate=0.05, momentum=0.5, num_epochs=5,
do_rao_blackwell=True)
def _make_sequences() -> _np_cpu.ndarray:
"""Binary sequences: each row is one sequence of T frames."""
rng = _np_cpu.random.RandomState(0)
rng = _np_cpu.random.RandomState(0)
base = (rng.rand(NUM_SEQ, SENSORY_SIZE) > 0.5).astype(_np_cpu.float64)
seqs = _np_cpu.stack([base] * T, axis=1) # (NUM_SEQ, T, SENSORY_SIZE)
return seqs
return _np_cpu.stack([base] * T, axis=1) # (NUM_SEQ, T, SENSORY_SIZE)
def test_rnn_single_layer():
def test_rnn_shared():
"""Shared-weights mode: one entity reused at every time step."""
seqs = _make_sequences()
rnn = StackRnn("test_rnn", WORK_DIR)
layer = StackRnn.make_layer(
"layer0", SENSORY_SIZE, H_SIZE,
EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False),
TrainingParams(learning_rate=0.05, momentum=0.5, num_epochs=20,
mini_batch_size=0, do_rao_blackwell=True),
)
rnn.append(layer)
rnn = StackRnn("test_rnn_shared", WORK_DIR)
rnn.append(StackRnn.make_layer("layer0", SENSORY_SIZE, H_SIZE,
EntityParams(), _PARAMS))
rnn.state_init(0.01)
# Confirm entity shape
assert rnn.sensory_size() == SENSORY_SIZE, "sensory_size mismatch"
assert rnn.h_size() == H_SIZE, "h_size mismatch"
assert rnn.is_shared
assert rnn.sensory_size() == SENSORY_SIZE
assert rnn.h_size() == H_SIZE
rnn.train(seqs)
# Inference: step through one sequence
rnn.reset(batch_size=1)
seq0 = seqs[0] # (T, SENSORY_SIZE) numpy
for t in range(T):
x_t = np.array(seq0[t][None, :]) # (1, SENSORY_SIZE) on device
h = rnn.step(x_t)
assert h.shape == (1, H_SIZE), f"step output shape wrong at t={t}"
h = rnn.step(np.array(seqs[0, t][None, :]))
assert h.shape == (1, H_SIZE), f"bad shape at t={t}"
# Reconstruction from final hidden state
recon = rnn.reconstruct(h)
assert recon.shape == (1, SENSORY_SIZE), "reconstruct shape wrong"
assert recon.shape == (1, SENSORY_SIZE)
print(f"Original x_T : {convert(np.array(seq0[-1][None, :]))}")
print(f"Reconstructed: {convert(recon)}")
print("test_rnn_single_layer: [passed]")
print(f"Original : {convert(np.array(seqs[0, -1][None, :]))}")
print(f"Recon : {convert(recon)}")
print("test_rnn_shared: [passed]")
def test_rnn_two_layers():
"""Two-layer recurrent stack: layer 1 receives h_0 as its sensory input."""
H_SIZE_0, H_SIZE_1 = 6, 3
def test_rnn_unrolled():
"""Unrolled mode: T entities, one per time step, each with own weights."""
seqs = _make_sequences()
rnn = StackRnn("test_rnn2", WORK_DIR)
rnn.append(StackRnn.make_layer(
"layer0", SENSORY_SIZE, H_SIZE_0,
EntityParams(),
TrainingParams(learning_rate=0.05, num_epochs=10),
))
rnn.append(StackRnn.make_layer(
"layer1", H_SIZE_0, H_SIZE_1,
EntityParams(),
TrainingParams(learning_rate=0.05, num_epochs=10),
))
rnn = StackRnn("test_rnn_unrolled", WORK_DIR)
for layer in StackRnn.make_unrolled(T, SENSORY_SIZE, H_SIZE,
EntityParams(), _PARAMS):
rnn.append(layer)
rnn.state_init(0.01)
assert not rnn.is_shared
assert rnn.num_layers() == T
rnn.train(seqs)
# Inference — _t advances through all T entities
rnn.reset(batch_size=1)
for t in range(T):
x_t = np.array(seqs[0, t][None, :])
h = rnn.step(x_t)
h = rnn.step(np.array(seqs[0, t][None, :]))
assert h.shape == (1, H_SIZE), f"bad shape at t={t}"
assert rnn._t == t + 1
assert h.shape == (1, H_SIZE_1)
print("test_rnn_two_layers: [passed]")
recon = rnn.reconstruct(h)
assert recon.shape == (1, SENSORY_SIZE)
# next_entity wraps around modularly
rnn._t = T + 3
assert rnn.next_entity() is rnn.from_index(3).entity
print("test_rnn_unrolled: [passed]")
def test_rnn_save_load():
seqs = _make_sequences()
rnn = StackRnn("test_rnn_sl", WORK_DIR)
rnn.append(StackRnn.make_layer(
"layer0", SENSORY_SIZE, H_SIZE,
EntityParams(),
TrainingParams(num_epochs=5),
))
rnn = StackRnn("test_rnn_sl", WORK_DIR)
for layer in StackRnn.make_unrolled(T, SENSORY_SIZE, H_SIZE,
EntityParams(), _PARAMS):
rnn.append(layer)
rnn.state_init(0.01)
rnn.train(seqs)
rnn.state_save()
@@ -106,7 +93,7 @@ def test_rnn_save_load():
if __name__ == "__main__":
test_rnn_single_layer()
test_rnn_two_layers()
test_rnn_shared()
test_rnn_unrolled()
test_rnn_save_load()
print("All RNN tests passed.")