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

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:
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
+27 -18
View File
@@ -38,26 +38,31 @@ of a **context** vector (the previous hidden state) and the current sensory inpu
(next time step) (predict char)
```
### Temporal unrolling
### Temporal unrolling (unrolled / own-weights mode)
```
x_{t-1} x_t x_{t+1}
┌──────┴──────┐ ──────┴──────┐ ──────┴──────┐
│ ctx │x_{t-1} │ ctx │ x_t │ │ ctx │x_{t+1}│
│ ╠═══════╣ │ ╠═══════ │ ╠═══════╣
│ W (shared) │ W (shared) │ W (shared)
│ ╠═══════╣ │ ╠═══════ │ ╠═══════╣
hidden_t-1 hidden_t │ │ hidden_t+1
└───────┬───────┘ └───────┬─────── └───────┬───────┘
context_t-1 │ context_t context_t+1
└────────────────►┘ ────────────────►┘ ──────► ...
x_{t-1} x_t x_{t+1}
┌──────┴──────┐ ┌───────┴──────┐ ┌───────┴──────┐
│ ctx │x_{t-1} │ ctx │ x_t │ │ ctx │x_{t+1}│
│ ╠═══════╣ │ ╠════════╣ │ ╠═══════
│ W_{t-1} │ │ W_t │ │ W_{t+1}
│ ╠═══════╣ │ ╠════════╣ │ ╠═══════
│ hidden_{t-1} │ │ hidden_t │ hidden_{t+1}
└───────┬───────┘ └───────┬────────┘ └───────┬───────
│ context_{t-1} │ context_t context_{t+1}
└────────────────►┘ ──────────────────►┘ ──────► ...
```
Weights **W**, **b_v**, **b_h** are shared across all time steps — the same RBM
processes every character. This is the RTRBM concatenation variant: instead of
modulating the hidden biases (Sutskever & Hinton, 2007), the previous hidden
state is directly concatenated to the visible layer.
Each time step has its **own weight matrix** W_t, b_v_t, b_h_t — the model is
**unrolled**: N positions in a sequence → N separate RBMs. This lets each
position specialise for the statistical patterns that occur at that offset in a
sequence. The sequence length T must be fixed and equal to N at training time;
during generation, position indices wrap modulo N.
For comparison, the shared-weights variant (1 entity, `make_layer`) uses a
single W across all time steps — the RTRBM concatenation variant of Sutskever &
Hinton (2007).
---
@@ -78,7 +83,8 @@ All hyperparameters live in the **Build model** cell of `moby_rnn.ipynb`:
| `SEED` | `"Call me Ishmael."` | Seed text for free generation |
| `SEQ_IDX` | 42 | Sequence index for the hidden-state trace plot |
Model size with defaults: **(512 + 85) × 512 = 306,176 parameters**.
Model size with defaults: **(CONTEXT_SIZE + 85) × CONTEXT_SIZE × T parameters**
e.g. T=100, CONTEXT_SIZE=256 → 100 × 27,136 = **2,713,600 parameters** total.
---
@@ -107,7 +113,10 @@ continue training from the last saved state — `model.state_load()` in cell 4
picks it up automatically.
```
results/moby_rnn-0-state.npz ← single-layer checkpoint
results/moby_rnn-0-state.npz ← RBM at position 0
results/moby_rnn-1-state.npz ← RBM at position 1
...
results/moby_rnn-99-state.npz ← RBM at position T-1
```
---
+7 -54
View File
@@ -158,48 +158,9 @@
"start_time": "2026-05-31T09:38:49.498420274Z"
}
},
"source": [
"# ── Build model ────────────────────────────────────────────────────────────\n",
"# visible layer = [context (CONTEXT_SIZE) | x_t (vocab_size)]\n",
"# hidden layer = h_t (CONTEXT_SIZE)\n",
"rnn = StackRnn(PRJ_NAME, WORK_DIR)\n",
"rnn.append(StackRnn.make_layer(\n",
" \"layer0\",\n",
" sensory_size=vocab_size,\n",
" h_size=CONTEXT_SIZE,\n",
" entity_params=EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False),\n",
" training_params=TrainingParams(\n",
" learning_rate=0.005,\n",
" momentum=0.9,\n",
" num_epochs=200,\n",
" do_rao_blackwell=True,\n",
" l2_lambda=0.0001,\n",
" ),\n",
"))\n",
"\n",
"rnn.state_init(0.01)\n",
"rnn.state_load() # resumes from checkpoint if one exists\n",
"\n",
"e = rnn.from_index(0).entity\n",
"print(f\"Entity : {e.name}\")\n",
"print(f\"Visible : {rnn.h_size()} (context) + {rnn.sensory_size()} (vocab) = {e.shape[0]}\")\n",
"print(f\"Hidden : {rnn.h_size()}\")\n",
"print(f\"Parameters : {e.shape[0] * e.shape[1]:,}\")"
],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"results/moby_rnn-0-state.npz loaded successfully!\n",
"Entity : Entity-213x128\n",
"Visible : 128 (context) + 85 (vocab) = 213\n",
"Hidden : 128\n",
"Parameters : 27,264\n"
]
}
],
"execution_count": 58
"source": "# ── Build model ────────────────────────────────────────────────────────────\n# Unrolled mode: T separate RBMs, one per position in the sequence.\n# Each RBM_t has its own W_t, b_v_t, b_h_t.\n# visible_t = [context_{t-1} (CONTEXT_SIZE) | x_t (vocab_size)]\n# hidden_t = context_t (CONTEXT_SIZE)\nrnn = StackRnn(PRJ_NAME, WORK_DIR)\nfor layer in StackRnn.make_unrolled(\n T, vocab_size, CONTEXT_SIZE,\n entity_params=EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False),\n training_params=TrainingParams(\n learning_rate=0.005,\n momentum=0.9,\n num_epochs=200,\n do_rao_blackwell=True,\n l2_lambda=0.0001,\n ),\n):\n rnn.append(layer)\n\nrnn.state_init(0.01)\nrnn.state_load() # resumes from checkpoint if one exists\n\ne = rnn.from_index(0).entity\nprint(f\"Mode : unrolled ({rnn.num_layers()} layers, own weights per position)\")\nprint(f\"Visible : {rnn.h_size()} (context) + {rnn.sensory_size()} (vocab) = {e.shape[0]}\")\nprint(f\"Hidden : {rnn.h_size()}\")\nprint(f\"Params/layer : {e.shape[0] * e.shape[1]:,}\")\nprint(f\"Total params : {e.shape[0] * e.shape[1] * rnn.num_layers():,}\")",
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
@@ -349,17 +310,9 @@
"start_time": "2026-05-31T09:39:43.947416428Z"
}
},
"source": "# ── Generation helpers ─────────────────────────────────────────────────────\n\ndef predict_next(rnn, context, n_gibbs=10, temperature=1.0):\n \"\"\"Return a character probability distribution conditioned on context.\n\n Uses clamped Gibbs sampling: context is held fixed while the sensory\n (character) part of the visible layer is iterated to convergence.\n \"\"\"\n entity = rnn.from_index(0).entity\n h_sz = rnn.h_size()\n s_sz = rnn.sensory_size()\n\n x_init = (np.random.rand(1, s_sz) > 0.5).astype(float)\n visible = np.concatenate([context, x_init], axis=1)\n\n for _ in range(n_gibbs):\n h = entity.forward(visible)\n visible = entity.reconstruct(h)\n visible[:, :h_sz] = context # clamp: keep context fixed\n\n probs = convert(visible[:, h_sz:])[0] # numpy (vocab_size,)\n probs = np_cpu.power(np_cpu.clip(probs, 1e-10, 1.0), 1.0 / temperature)\n probs /= probs.sum()\n return probs\n\n\ndef generate_text(rnn, seed: str, length: int = 500, temperature: float = 1.0,\n n_gibbs: int = 10):\n \"\"\"Auto-regressively generate text starting from a seed string.\"\"\"\n rnn.reset(batch_size=1)\n h = np.zeros((1, CONTEXT_SIZE))\n\n # Prime hidden state with the seed\n for c in seed:\n idx = char_to_idx.get(c, 0)\n x = np.zeros((1, vocab_size))\n x[0, idx] = 1.0\n h = rnn.step(x)\n\n generated = seed\n for _ in range(length):\n probs = predict_next(rnn, h.copy(), n_gibbs=n_gibbs, temperature=temperature)\n idx = int(np_cpu.random.choice(vocab_size, p=probs))\n c = idx_to_char[idx]\n generated += c\n\n x = np.zeros((1, vocab_size))\n x[0, idx] = 1.0\n h = rnn.step(x)\n\n return generated\n\nprint(\"Helpers defined.\")",
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Helpers defined.\n"
]
}
],
"execution_count": 60
"source": "# ── Generation helpers ─────────────────────────────────────────────────────\n\ndef predict_next(rnn, context, n_gibbs=10, temperature=1.0):\n \"\"\"Return a character probability distribution conditioned on context.\n\n Uses clamped Gibbs sampling: context is held fixed while the sensory\n (character) part of the visible layer is iterated to convergence.\n Uses rnn.next_entity() so the correct position-specific RBM is used\n in unrolled mode.\n \"\"\"\n entity = rnn.next_entity()\n h_sz = rnn.h_size()\n s_sz = rnn.sensory_size()\n\n x_init = (np.random.rand(1, s_sz) > 0.5).astype(float)\n visible = np.concatenate([context, x_init], axis=1)\n\n for _ in range(n_gibbs):\n h = entity.forward(visible)\n visible = entity.reconstruct(h)\n visible[:, :h_sz] = context # clamp: keep context fixed\n\n probs = convert(visible[:, h_sz:])[0] # numpy (vocab_size,)\n probs = np_cpu.power(np_cpu.clip(probs, 1e-10, 1.0), 1.0 / temperature)\n probs /= probs.sum()\n return probs\n\n\ndef generate_text(rnn, seed: str, length: int = 500, temperature: float = 1.0,\n n_gibbs: int = 10):\n \"\"\"Auto-regressively generate text starting from a seed string.\"\"\"\n rnn.reset(batch_size=1)\n h = np.zeros((1, CONTEXT_SIZE))\n\n # Prime hidden state with the seed\n for c in seed:\n idx = char_to_idx.get(c, 0)\n x = np.zeros((1, vocab_size))\n x[0, idx] = 1.0\n h = rnn.step(x)\n\n generated = seed\n for _ in range(length):\n probs = predict_next(rnn, h.copy(), n_gibbs=n_gibbs, temperature=temperature)\n idx = int(np_cpu.random.choice(vocab_size, p=probs))\n c = idx_to_char[idx]\n generated += c\n\n x = np.zeros((1, vocab_size))\n x[0, idx] = 1.0\n h = rnn.step(x)\n\n return generated\n\nprint(\"Helpers defined.\")",
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
@@ -523,4 +476,4 @@
"execution_count": 65
}
]
}
}
+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": ""},
})
+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.")