diff --git a/src/rbm/rbm.py b/src/rbm/rbm.py index 379694a..e8dafa7 100644 --- a/src/rbm/rbm.py +++ b/src/rbm/rbm.py @@ -45,7 +45,7 @@ def read_armadillo(filename: str) -> np.ndarray: def main(): # Create stack from project file - stack = StackFactory.from_file("/home/jens/work/repos/Rbm/test.prj") + stack = StackFactory.from_file("/home/jens/work/repos/Rbm/test.prj", work_dir="../../results") # Load state stack.state_load() diff --git a/src/rbm/stack.py b/src/rbm/stack.py index 8724914..ebabf99 100644 --- a/src/rbm/stack.py +++ b/src/rbm/stack.py @@ -1,3 +1,4 @@ +import os from enum import Enum from layer import Layer @@ -9,10 +10,12 @@ class StackException(Exception): pass class Stack: - def __init__(self, stack_type: StackType, name: str): + def __init__(self, stack_type: StackType, name: str, work_dir: str = '.'): self.stack_type = stack_type self.name = name + self.work_dir = work_dir self.layers: list[Layer] = [] + os.makedirs(self.work_dir, exist_ok=True) def num_layers(self): return len(self.layers) @@ -44,8 +47,10 @@ class Stack: def state_save(self): for index, layer in enumerate(self.layers): - layer.save(f"{self.name}-{index}-state.npz") + filepath = os.path.join(self.work_dir, f"{self.name}-{index}-state.npz") + layer.save(filepath) def state_load(self): for index, layer in enumerate(self.layers): - layer.load(f"{self.name}-{index}-state.npz") \ No newline at end of file + filepath = os.path.join(self.work_dir, f"{self.name}-{index}-state.npz") + layer.load(filepath) diff --git a/src/rbm/stack_deep.py b/src/rbm/stack_deep.py index 8cbec59..3826358 100644 --- a/src/rbm/stack_deep.py +++ b/src/rbm/stack_deep.py @@ -4,8 +4,8 @@ from cd_train import cd_jens from stack import Stack, StackType class StackDeep(Stack): - def __init__(self, name: str): - Stack.__init__(self, StackType.Deep, name) + def __init__(self, name: str, work_dir: str = '.'): + Stack.__init__(self, StackType.Deep, name, work_dir) def batch_from(self, batch: np.ndarray, from_layer_id: int = 0): _batch = np.matrix.copy(batch) diff --git a/src/rbm/stack_factory.py b/src/rbm/stack_factory.py index 2ae7275..bf9aed3 100644 --- a/src/rbm/stack_factory.py +++ b/src/rbm/stack_factory.py @@ -10,7 +10,7 @@ from stack_rnn import StackRnn class StackFactory: @classmethod - def from_dict(cls, project: dict, layer_constructor: Callable = None) -> StackDeep|StackRnn: + def from_dict(cls, project: dict, work_dir: str = ".", layer_constructor: Callable = None) -> StackDeep|StackRnn: name = project["stack"]["name"] layers = project["stack"]["layers"] @@ -21,9 +21,9 @@ class StackFactory: obj = None if stack_type == StackType.Deep: - obj = StackDeep(name) + obj = StackDeep(name, work_dir) if stack_type == StackType.Rnn: - obj = StackRnn(name) + obj = StackRnn(name, work_dir) if obj is None: return obj @@ -56,16 +56,16 @@ class StackFactory: return obj @classmethod - def from_file(cls, filename: str, layer_constructor: Callable = None) -> StackDeep|StackRnn: + def from_file(cls, filename: str, work_dir: str = ".", layer_constructor: Callable = None) -> StackDeep|StackRnn: obj = None with open(filename, "r") as fp: prj = json.load(fp) - obj = StackFactory.from_dict(prj) + obj = StackFactory.from_dict(prj, work_dir) return obj if __name__ == "__main__": - stack = StackFactory.from_file("/home/jens/work/repos/Rbm/test.prj") + stack = StackFactory.from_file("/home/jens/work/repos/Rbm/test.prj", work_dir="../../results") print("Test: [passed]") diff --git a/src/rbm/stack_rnn.py b/src/rbm/stack_rnn.py index 81b6d79..7061026 100644 --- a/src/rbm/stack_rnn.py +++ b/src/rbm/stack_rnn.py @@ -1,5 +1,5 @@ from stack import Stack, StackType class StackRnn(Stack): - def __init__(self, name: str): - Stack.__init__(self, StackType.Rnn, name) + def __init__(self, name: str, work_dir: str = '.'): + Stack.__init__(self, StackType.Rnn, name, work_dir)