added model

src/rbm/model.py

@@ -0,0 +1,45 @@
+import os
+from abc import ABC, abstractmethod
+
+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
+
+	def objects(self, obj_type: type = Entity) -> list[Entity]:
+		obj_list: list[type[obj_type]] = []
+		for name, value in self.__dict__.items():
+			if isinstance(value, obj_type):
+				obj_list.append(value)
+		return obj_list
+
+	def train(self, batch: Mat, params: TrainingParams):
+		_batch = np.copy(batch)
+		for entity in self.objects(Entity):
+			train(entity, _batch, params, Status(), cd_jens)
+			_batch = entity(_batch)
+
+	@abstractmethod
+	def forward(self, x: Mat) -> Mat:
+		pass
+
+	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)
+
+	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
+
+