From 3698c85ed43e788b6715b6bba5a0d176e149e376 Mon Sep 17 00:00:00 2001 From: Jens Ahrensfeld Date: Sun, 31 May 2026 11:52:01 +0200 Subject: [PATCH] [StackRnn] - add unrolled (own-weights) mode; each position gets its own RBM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- README_moby_rnn.md | 45 +++++---- moby_rnn.ipynb | 61 ++----------- src/rbm/stack_rnn.py | 208 +++++++++++++++++++++++++++++------------- src/tests/test_rnn.py | 115 +++++++++++------------ 4 files changed, 229 insertions(+), 200 deletions(-) diff --git a/README_moby_rnn.md b/README_moby_rnn.md index 095a3d2..e19f2fe 100644 --- a/README_moby_rnn.md +++ b/README_moby_rnn.md @@ -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 ``` --- diff --git a/moby_rnn.ipynb b/moby_rnn.ipynb index 7760f45..73d618e 100644 --- a/moby_rnn.ipynb +++ b/moby_rnn.ipynb @@ -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 } ] -} +} \ No newline at end of file diff --git a/src/rbm/stack_rnn.py b/src/rbm/stack_rnn.py index 835bce3..3d0ea19 100644 --- a/src/rbm/stack_rnn.py +++ b/src/rbm/stack_rnn.py @@ -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": ""}, + }) diff --git a/src/tests/test_rnn.py b/src/tests/test_rnn.py index 530059f..313761f 100644 --- a/src/tests/test_rnn.py +++ b/src/tests/test_rnn.py @@ -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.")