93 lines
3.1 KiB
Python
93 lines
3.1 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
|
|
from rbm.train import train, Status
|
|
from label.label import Label
|
|
|
|
class TestModel(Model):
|
|
def __init__(self, name: str, work_dir: str = '.'):
|
|
super().__init__(name, work_dir)
|
|
# self.unit_image = Entity((96*96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False, num_gibbs_samples=3), TrainingParams(learning_rate=0.00001, momentum=0.9, do_rao_blackwell=False, num_epochs=1000))
|
|
self.unit_image = Entity((96 * 96, 16), EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False))
|
|
self.unit_combine = Entity((32, 64), EntityParams(), TrainingParams(learning_rate=0.01, momentum=0.9, do_rao_blackwell=True, num_epochs=1000))
|
|
|
|
def forward2(self, images: Mat, labels: Mat):
|
|
h = self.unit_image.forward(images)
|
|
h = np.ndarray.flatten(h)
|
|
v_combine = np.concat((h, labels), axis=0)
|
|
h_res = self.unit_combine.forward(v_combine)
|
|
return h_res
|
|
|
|
def reconstruct(self, h: Mat):
|
|
v_combine = self.unit_combine.reconstruct(h)
|
|
v_image = self.unit_image.reconstruct(np.transpose(np.transpose(v_combine)[0:self.unit_image.shape[1]]))
|
|
v_label = np.transpose(np.transpose(v_combine)[self.unit_image.shape[1]:self.unit_image.shape[1]+self.unit_image.shape[1]])
|
|
return np.ndarray.flatten(v_image), np.ndarray.flatten(v_label)
|
|
|
|
def train2(self, images: Mat, labels: Mat):
|
|
train(self.unit_image, images, Status())
|
|
h11 = self.unit_image.forward(images)
|
|
x2 = np.concat((h11, labels), axis=1)
|
|
train(self.unit_combine, x2, Status())
|
|
|
|
if __name__ == "__main__":
|
|
work_dir = "results"
|
|
prj_name = "norb_small_16h_v2"
|
|
prj_root = "/home/jens/work/repos/Rbm"
|
|
|
|
# Create model
|
|
model = TestModel("norb_small_16h_v2", "results")
|
|
|
|
# Init state
|
|
model.init(0.01)
|
|
|
|
# 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
|
|
# Create Labeler
|
|
enc = Label(20, 16, Label.EncodingType.OneHot, work_dir)
|
|
|
|
# Load
|
|
enc.fitter.load()
|
|
|
|
# Encode test labels
|
|
train_labels = np.array([1, 2, 3, 4, 5])
|
|
enc.fit(train_labels)
|
|
|
|
encoded, _ = enc.encode(train_labels)
|
|
|
|
model.train2(train_batch, encoded)
|
|
|
|
# save state
|
|
model.save()
|
|
|
|
cmap = 'Grays'
|
|
fig, axes = plt.subplots(3, len(train_batch), figsize=(12, 3))
|
|
for index, image in enumerate(train_batch):
|
|
label_zero = np.zeros(shape=16)
|
|
image_zero = np.zeros(shape=96*96)
|
|
image_reconst, labels_reconst = model.reconstruct(model.forward2(image, label_zero))
|
|
# image_reconst, labels_reconst = model.reconstruct(model.forward2(image_zero, encoded[index, :]))
|
|
img = 1-2*(image_reconst + 0.5)
|
|
img = np.reshape(img, (96, 96))
|
|
axes[0][index].imshow(img, cmap=cmap)
|
|
axes[0][index].axis('off')
|
|
axes[1][index].imshow(np.reshape(labels_reconst, (4,4)), cmap=cmap)
|
|
axes[1][index].axis('off')
|
|
axes[2][index].imshow(np.reshape(encoded[index, :], (4,4)), cmap=cmap)
|
|
axes[2][index].axis('off')
|
|
|
|
plt.show()
|
|
|
|
|
|
|