diff --git a/src/rbm/model.py b/src/rbm/model.py index be3fff9..e8bfd5b 100644 --- a/src/rbm/model.py +++ b/src/rbm/model.py @@ -20,12 +20,12 @@ class Model(ABC): obj_list.append(value) return obj_list - def train(self, batch: Mat): + def train(self, batch: Mat, status: Status = None): entities = self.objects(Entity) _batch = np.array(batch) # np.array() converts numpy→CuPy; np.copy() does not for entity in entities: if entity.enable_training: - train(entity, _batch, Status()) + train(entity, _batch, status if status is not None else Status()) _batch = entity.forward(_batch) def forward(self, x: Mat) -> Mat: diff --git a/src/rbm/status.py b/src/rbm/status.py index ee7cd76..4d22c57 100644 --- a/src/rbm/status.py +++ b/src/rbm/status.py @@ -36,4 +36,17 @@ class Status: def on_report(self, entity: Entity, status: dict) -> bool: Status.print_status(entity, status) + return True + + +class CheckpointStatus(Status): + """Prints progress and saves model state via save_fn at every report interval.""" + + def __init__(self, save_fn, update_interval: int = 10): + super().__init__(update_interval) + self.save_fn = save_fn + + def on_report(self, entity: Entity, status: dict) -> bool: + Status.print_status(entity, status) + self.save_fn() return True \ No newline at end of file diff --git a/src/tests/test_faces_sub_image.py b/src/tests/test_faces_sub_image.py index 86c01fd..8f99f07 100644 --- a/src/tests/test_faces_sub_image.py +++ b/src/tests/test_faces_sub_image.py @@ -7,6 +7,7 @@ from rbm.model import Model from rbm.entity import Entity, EntityParams, TrainingParams from rbm.image import SubImage, normalize from rbm.matrix import Mat, np, convert +from rbm.status import CheckpointStatus DATA_DIR = '/media/jens/cifs/bilder/MachineVision/Caltech_WebFaces/' PATCH = 32 @@ -214,7 +215,7 @@ if __name__ == '__main__': print(f'Loading {args.n_images} training images (stride={STRIDE})...') train_patches = load_patches(DATA_DIR, args.n_images) print(f'Training on {train_patches.shape[0]} patches ({N_VIS}-dim each)') - model.train(train_patches) + model.train(train_patches, status=CheckpointStatus(save_fn=model.save)) model.save() # Show learned weight filters