47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
import os
|
|
from abc import ABC, abstractmethod
|
|
|
|
from .entity import Entity
|
|
from .train import train
|
|
from .matrix import Mat, np
|
|
from .status import Status
|
|
|
|
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]] = []
|
|
for name, value in self.__dict__.items():
|
|
if isinstance(value, obj_type):
|
|
obj_list.append(value)
|
|
return obj_list
|
|
|
|
def train(self, batch: Mat):
|
|
entities = self.objects(Entity)
|
|
_batch = np.copy(batch)
|
|
for entity in entities:
|
|
if entity.enable_training:
|
|
train(entity, _batch, Status())
|
|
_batch = entity.forward(_batch)
|
|
|
|
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.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")
|
|
entity.state.load(filepath)
|
|
|
|
def init(self, std: float):
|
|
for index, entity in enumerate(self.objects(Entity)):
|
|
entity.state.init(std=std)
|