diff --git a/src/rbm/entity.py b/src/rbm/entity.py index 2f75592..77e5fb2 100644 --- a/src/rbm/entity.py +++ b/src/rbm/entity.py @@ -2,10 +2,11 @@ from .state import RbmState from .matrix import prob, Mat class EntityParams: - def __init__(self, do_gaussian_visible: bool = False, do_gaussian_hidden: bool = False): + def __init__(self, do_gaussian_visible: bool = False, do_gaussian_hidden: bool = False, num_gibbs_samples: int = 1): # Entity parameters self.do_gaussian_visible = do_gaussian_visible self.do_gaussian_hidden = do_gaussian_hidden + self.num_gibbs_samples = num_gibbs_samples @classmethod def from_dict(cls, params: dict): @@ -14,6 +15,8 @@ class EntityParams: obj.do_gaussian_visible = params["doGaussianVisible"] if "doGaussianHidden" in params: obj.do_gaussian_hidden = params["doGaussianHidden"] + if "num_gibbs_samples" in params: + obj.num_gibbs_samples = params["num_gibbs_samples"] return obj @@ -44,7 +47,8 @@ class Entity: self.state.b_h += grad.b_h self.state.w_hv += grad.w_hv - def forward(self, v: Mat, num_gibbs: int = 1) -> Mat: + def forward(self, v: Mat, num_gibbs: int = 0) -> Mat: + num_gibbs = self.params.num_gibbs_samples if num_gibbs == 0 else num_gibbs h = self._v_to_ph(v) for i in range(num_gibbs-1): h = self._h_to_pv(h) @@ -52,7 +56,8 @@ class Entity: return h - def reconstruct(self, h: Mat, num_gibbs: int = 1) -> Mat: + def reconstruct(self, h: Mat, num_gibbs: int = 0) -> Mat: + num_gibbs = self.params.num_gibbs_samples if num_gibbs == 0 else num_gibbs v = self._h_to_pv(h) for i in range(num_gibbs-1): v = self._v_to_ph(v) diff --git a/src/rbm/layer.py b/src/rbm/layer.py index 327b26a..aa153c0 100644 --- a/src/rbm/layer.py +++ b/src/rbm/layer.py @@ -13,13 +13,10 @@ class Layer: self.entity.state.init(mu=0, std=std) def save(self, filename: str = None): - self.entity.state.to_file(filename) + self.entity.state.save(filename) def load(self, filename: str = None): - state = RbmState.from_file(filename) - if state is not None: - self.entity.state = state - + self.entity.state.load(filename) diff --git a/src/rbm/model.py b/src/rbm/model.py index ae5b333..b319f68 100644 --- a/src/rbm/model.py +++ b/src/rbm/model.py @@ -5,13 +5,13 @@ from .entity import Entity from .train import train, TrainingParams, cd_jens from .matrix import Mat, np from .status import Status -from .state import RbmState class Model(ABC): known_classes = [Entity] def __init__(self, name: str = "myStack", work_dir: str = "."): self.name = name self.work_dir = work_dir + os.makedirs(self.work_dir, exist_ok=True) def objects(self, obj_type: type = Entity) -> list[Entity]: obj_list: list[type[obj_type]] = [] @@ -24,7 +24,7 @@ class Model(ABC): _batch = np.copy(batch) for entity in self.objects(Entity): train(entity, _batch, params, Status(), cd_jens) - _batch = entity(_batch) + _batch = entity.forward(_batch, num_gibbs=params.num_gibbs_samples) @abstractmethod def forward(self, x: Mat) -> Mat: @@ -33,13 +33,13 @@ class Model(ABC): def save(self): for index, entity in enumerate(self.objects(Entity)): filepath = os.path.join(self.work_dir, f"{self.name}-{index}-state.npz") - entity.state.to_file(filepath) + entity.state.save(filepath) def load(self): for index, entity in enumerate(self.objects(Entity)): filepath = os.path.join(self.work_dir, f"{self.name}-{index}-state.npz") - state = RbmState.from_file(filepath) - if state is not None: - entity.state = state - + entity.state.load(filepath) + def init(self, std: float): + for index, entity in enumerate(self.objects(Entity)): + entity.state.init(std=std) diff --git a/src/rbm/state.py b/src/rbm/state.py index 7857da3..1dc557f 100644 --- a/src/rbm/state.py +++ b/src/rbm/state.py @@ -36,7 +36,22 @@ class RbmState: def h_to_v(self, hidden: Mat) -> Mat: return np.dot(hidden, np.transpose(self.w_hv)) + self.b_v - def to_file(self, filename: str): + def load(self, filename: str): + try: + with np.load(filename) as X: + w_hv, b_v, b_h = [X[i] for i in ('whv', 'bv', 'bh')] + if w_hv.shape == self.w_hv.shape and b_v.shape == self.b_v.shape and b_h.shape == self.b_h.shape: + self.w_hv = w_hv + self.b_v = b_v + self.b_h = b_h + print(f"{filename} loaded successfully!") + except FileNotFoundError: + pass + except KeyError: + pass + + + def save(self, filename: str): np.savez(filename, whv=self.w_hv, bv=self.b_v, bh=self.b_h) print(f"{filename} saved successfully!") diff --git a/src/tests/test_model.py b/src/tests/test_model.py index 23e06eb..1abe717 100644 --- a/src/tests/test_model.py +++ b/src/tests/test_model.py @@ -5,36 +5,46 @@ from rbm.entity import Entity, EntityParams from rbm.matrix import Mat, np from rbm.train import TrainingParams -class DeepStack(Model): - def __init__(self, name: str = "myStack"): - super().__init__(name) - self.unit1 = Entity((32*32*3, 256), EntityParams(do_gaussian_visible=True)) -# self.unit2 = Entity((16, 24), EntityParams()) -# self.unit3 = Entity((24, 10), EntityParams()) +class TestModel(Model): + def __init__(self, name: str, work_dir: str = '.'): + super().__init__(name, work_dir) + self.unit1 = Entity((16, 64), EntityParams()) + self.unit2 = Entity((64, 16), EntityParams()) + self.unit3 = Entity((16, 64), EntityParams()) def forward(self, x: Mat): - x = self.unit1(x) -# x = self.unit2(x) -# x = self.unit3(x) + x = self.unit1.forward(x) + x = self.unit2.forward(x) + x = self.unit3.forward(x) return x + def backward(self, x: Mat): + x = self.unit3.reconstruct(x) + x = self.unit2.reconstruct(x) + x = self.unit1.reconstruct(x) + return x if __name__ == "__main__": # Create model - model = DeepStack("myStack") + model = TestModel("TestModel", "results") + + # Init state + model.init(0.01) # load state model.load() # create batch - batch = (np.random.rand(16, 32*32*3) > 0.5).astype(np.float64) + batch = (np.random.rand(64, 16) > 0.5).astype(np.float64) # Train - model.train(batch, TrainingParams(learning_rate=0.1, momentum=0.9, do_rao_blackwell=True, num_epochs=1000)) + model.train(batch, TrainingParams(learning_rate=0.01, momentum=0.9, do_rao_blackwell=True, num_epochs=1000, num_gibbs_samples=3)) # save state model.save() - for pat in batch: - model.forward(pat) + for inp in batch: + out = model.backward(model.forward(inp)) + print(f"Pattern: {inp} -> {out}") +