Files
pyRBM/tests/test_norbs.py
jensandClaude Sonnet 4.6 3edc124a5c [refactor] move tests to tests/, add pytest functions and main() entrypoints
- Moved src/tests/ → tests/
- Added test_* functions with assertions to script-style test files
- Added main() to each so IDEs offer it as a separate run target from pytest
- Fixed cupy_test.py: remove spurious x_gpu += x_cpu, fix duplicate xlabel

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-02 21:32:58 +02:00

72 lines
1.7 KiB
Python

import os
import matplotlib.pyplot as plt
from model.model import Model
from rbm.entity import Entity, EntityParams, TrainingParams
from rbm.matrix import Mat, np, read_armadillo
DO_HIDDEN_GAUSSIAN = True
class TestModel(Model):
def __init__(self, name: str, work_dir: str = '.'):
super().__init__(name, work_dir)
if DO_HIDDEN_GAUSSIAN:
# Hidden gaussian
self.unit1 = Entity((96*96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=True),
TrainingParams(learning_rate=0.00001, momentum=0.9, num_epochs=1000))
else:
# Hidden binary
self.unit1 = Entity((96 * 96, 333), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=1000, l2_lambda=0.001))
def forward(self, x: Mat):
x = self.unit1.forward(x)
return x
def reconstruct(self, x: Mat):
x = self.unit1.reconstruct(x)
return x
if __name__ == "__main__":
work_dir = "results"
prj_name = "norb_small_16h_v2"
prj_root = "/home/jens/work/repos/Rbm"
# Create model
model = TestModel(prj_name, "results")
# Init state
if DO_HIDDEN_GAUSSIAN:
model.init(0.01)
else:
model.init(0.1)
# load state
# model.load()
# Load train data
train_batch = read_armadillo(os.path.join(prj_root, f"{prj_name}.training.dat"))
# Load test data
test_batch = read_armadillo(os.path.join(prj_root, f"{prj_name}.test.dat"))
# Train
model.train(train_batch)
# save state
model.save()
fig, axes = plt.subplots(1, len(test_batch), figsize=(12, 3))
for index, inp in enumerate(test_batch):
out_normalized = model.reconstruct(model.forward(inp))
img = 2*(out_normalized + 0.5)
img = np.reshape(img, (96, 96))
axes[index].imshow(np.asnumpy(img))
axes[index].axis('off')
plt.show()