67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
import os
|
|
import matplotlib.pyplot as plt
|
|
|
|
from rbm.model import Model
|
|
from rbm.entity import Entity, EntityParams, TrainingParams
|
|
from rbm.matrix import Mat, np, read_armadillo
|
|
|
|
class TestModel(Model):
|
|
def __init__(self, name: str, work_dir: str = '.', do_gaussian_hidden=False):
|
|
super().__init__(name, work_dir)
|
|
|
|
if do_gaussian_hidden:
|
|
# Hidden gaussian
|
|
self.unit1 = Entity((96*96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=True),
|
|
TrainingParams(learning_rate=0.000001, 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.01))
|
|
|
|
def forward(self, x: Mat):
|
|
x = self.unit1.forward(x)
|
|
return x
|
|
|
|
def backward(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
|
|
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.backward(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()
|
|
|
|
|
|
|