diff --git a/deeplay/applications/application.py b/deeplay/applications/application.py index 17fd191e..790e4d3d 100644 --- a/deeplay/applications/application.py +++ b/deeplay/applications/application.py @@ -124,9 +124,9 @@ def fit( ) history = LogHistory() - progressbar = RichProgressBar() + aux_callbacks = [history] - callbacks = callbacks + [history, progressbar] + callbacks = callbacks + aux_callbacks trainer = dl.Trainer(max_epochs=max_epochs, callbacks=callbacks, **kwargs) train_dataloader = torch.utils.data.DataLoader( @@ -243,12 +243,15 @@ def configure_optimizers(self): ) from e def training_step(self, batch, batch_idx): + x, y = self.train_preprocess(batch) y_hat = self(x) loss = self.compute_loss(y_hat, y) if not isinstance(loss, dict): loss = {"loss": loss} + assert "loss" in loss, "the output of compute_loss should contain a 'loss' key" + for name, v in loss.items(): self.log( f"train_{name}", @@ -263,7 +266,7 @@ def training_step(self, batch, batch_idx): "train", y_hat, y, on_step=True, on_epoch=True, prog_bar=True, logger=True ) - return sum(loss.values()) + return loss["loss"] def validation_step(self, batch, batch_idx): x, y = self.val_preprocess(batch) @@ -290,7 +293,7 @@ def validation_step(self, batch, batch_idx): prog_bar=True, logger=True, ) - return sum(loss.values()) + return loss["loss"] if "loss" in loss else 0 def test_step(self, batch, batch_idx): x, y = self.test_preprocess(batch) @@ -318,7 +321,7 @@ def test_step(self, batch, batch_idx): logger=True, ) - return sum(loss.values()) + return loss["loss"] if "loss" in loss else 0 def predict_step(self, batch, batch_idx, dataloader_idx=None): if isinstance(batch, (list, tuple)): @@ -356,12 +359,16 @@ def trainer(self, trainer): if module is self: continue try: - if hasattr(module, "trainer") and module.trainer is not trainer: + if isinstance(module, L.LightningModule) or hasattr(module, "trainer"): + module.trainer = trainer + except RuntimeError: # hasattr can raise RuntimeError if the module is not attached to a trainer if isinstance(module, L.LightningModule): + print("Battaching trainer to", module) module.trainer = trainer + print("Battached trainer to", module) @staticmethod def clone_metrics(metrics: T) -> T: diff --git a/deeplay/applications/detection/lodestar/transforms.py b/deeplay/applications/detection/lodestar/transforms.py index 157eea5c..6e794ceb 100644 --- a/deeplay/applications/detection/lodestar/transforms.py +++ b/deeplay/applications/detection/lodestar/transforms.py @@ -90,6 +90,15 @@ def _backward(x, angle, indices): mat2d[:, indices[1], indices[0]] = -torch.sin(-angle) mat2d[:, indices[0], indices[1]] = torch.sin(-angle) mat2d[:, indices[0], indices[0]] = torch.cos(-angle) - out = torch.matmul(x.unsqueeze(1), mat2d).squeeze(1) - - return out + + if len(x.size()) == 2: + # (B, C) -> (B, 1, C) + x = x.unsqueeze(1) + return torch.matmul(x, mat2d).squeeze(1) + + x_expanded = x.view(x.size(0), 1, x.size(1), -1) + y = torch.einsum("bijm,bjk->bikm", x_expanded, mat2d) + + return y.view(x.size()) + + \ No newline at end of file diff --git a/deeplay/module.py b/deeplay/module.py index 504ac5f0..53591da5 100644 --- a/deeplay/module.py +++ b/deeplay/module.py @@ -723,6 +723,39 @@ def replace(self, target: str, replacement: nn.Module): self._modules[target] = replacement + @after_init + def schedule(self, **schedulers): + for attr, scheduler in schedulers.items(): + setattr(self, attr, scheduler) + + @after_init + def schedule_linear( + self, + attr: str, + start_value: float, + end_value: float, + n_steps: int, + on_epoch: bool = False, + ): + from deeplay.schedulers import LinearScheduler + + setattr(self, attr, LinearScheduler(start_value, end_value, n_steps, on_epoch)) + + @after_init + def schedule_loglinear( + self, + attr: str, + start_value: float, + end_value: float, + n_steps: int, + on_epoch: bool = False, + ): + from deeplay.schedulers import LogLinearScheduler + + setattr( + self, attr, LogLinearScheduler(start_value, end_value, n_steps, on_epoch) + ) + @stateful def configure(self, *args: Any, **kwargs: Any): """ @@ -1335,6 +1368,14 @@ def __setattr__(self, name, value): # # ensure that logs are stored in the correct place # value.set_root_module(self.root_module) + def __getattr__(self, name): + from deeplay.schedulers import BaseScheduler + + x = super().__getattr__(name) + if self._has_built and isinstance(x, BaseScheduler): + return x.__get__(self, type(self)) + return x + @stateful def _set_submodule(self, name, module, tags): diff --git a/deeplay/schedulers/__init__.py b/deeplay/schedulers/__init__.py new file mode 100644 index 00000000..421de1b2 --- /dev/null +++ b/deeplay/schedulers/__init__.py @@ -0,0 +1,5 @@ +from .scheduler import BaseScheduler +from .linear import LinearScheduler +from .constant import ConstantScheduler +from .loglinear import LogLinearScheduler +from .sequence import SchedulerSequence diff --git a/deeplay/schedulers/constant.py b/deeplay/schedulers/constant.py new file mode 100644 index 00000000..87297929 --- /dev/null +++ b/deeplay/schedulers/constant.py @@ -0,0 +1,18 @@ +from . import BaseScheduler + + +class ConstantScheduler(BaseScheduler): + """Sheduler that returns constant value.""" + + def __init__(self, value, on_epoch=False): + super().__init__(on_epoch) + self.value = value + + def __call__(self, step): + return self.value + + def __repr__(self): + return f"{self.__class__.__name__}({self.value})" + + def __str__(self): + return repr(self) diff --git a/deeplay/schedulers/linear.py b/deeplay/schedulers/linear.py new file mode 100644 index 00000000..0cf65c49 --- /dev/null +++ b/deeplay/schedulers/linear.py @@ -0,0 +1,42 @@ +from .scheduler import BaseScheduler + + +class LinearScheduler(BaseScheduler): + """Scheduler that returns linearly changing value from start_value to end_value. + + For steps beyond n_steps, returns end_value. + For steps before 0, returns start_value. + + Parameters + ---------- + start_value : float + Initial value of the scheduler. + end_value : float + Final value of the scheduler. + n_steps : int + Number of steps to reach end_value. + on_epoch : bool + If True, the step is taken from the epoch counter of the trainer. + Otherwise, the step is taken from the global step counter of the trainer. + """ + + def __init__(self, start_value, end_value, n_steps, on_epoch=False): + super().__init__(on_epoch) + self.start_value = start_value + self.end_value = end_value + self.n_steps = n_steps + + def __call__(self, step): + if step < 0: + return self.start_value + if step >= self.n_steps: + return self.end_value + return ( + self.start_value + (self.end_value - self.start_value) * step / self.n_steps + ) + + def __repr__(self): + return f"{self.__class__.__name__}({self.start_value}, {self.end_value}, {self.n_steps})" + + def __str__(self): + return repr(self) diff --git a/deeplay/schedulers/loglinear.py b/deeplay/schedulers/loglinear.py new file mode 100644 index 00000000..ffc7a362 --- /dev/null +++ b/deeplay/schedulers/loglinear.py @@ -0,0 +1,33 @@ +from . import BaseScheduler +import numpy as np + + +class LogLinearScheduler(BaseScheduler): + """Scheduler that returns log-linearly changing value from start_value to end_value. + + For steps beyond n_steps, returns end_value.""" + + def __init__(self, start_value, end_value, n_steps, on_epoch=False): + super().__init__(on_epoch) + assert np.sign(start_value) == np.sign( + end_value + ), "Start and end values must have the same sign" + assert start_value != 0, "Start value must be non-zero" + assert end_value != 0, "End value must be non-zero" + assert n_steps > 0, "Number of steps must be greater than 0" + self.start_value = start_value + self.end_value = end_value + self.n_steps = n_steps + + def __call__(self, step): + if step >= self.n_steps: + return self.end_value + return self.start_value * (self.end_value / self.start_value) ** ( + step / self.n_steps + ) + + def __repr__(self): + return f"{self.__class__.__name__}({self.start_value}, {self.end_value}, {self.n_steps})" + + def __str__(self): + return repr(self) diff --git a/deeplay/schedulers/scheduler.py b/deeplay/schedulers/scheduler.py new file mode 100644 index 00000000..dafce25c --- /dev/null +++ b/deeplay/schedulers/scheduler.py @@ -0,0 +1,47 @@ +import lightning as L + +from deeplay.module import DeeplayModule +from deeplay.trainer import Trainer + + +class BaseScheduler(DeeplayModule, L.LightningModule): + """Base class for annealers.""" + + step: int + + def __init__(self, on_epoch=False): + super().__init__() + self.on_epoch = on_epoch + self._step = 0 + self._x = None + + def set_step(self, step): + self._step = step + self._x = self(step) + + def update(self): + current_step = self._step + + if self._trainer: + updated_step = ( + self.trainer.current_epoch + if self.on_epoch + else self.trainer.global_step + ) + else: + updated_step = self._step + + if updated_step != current_step or self._x is None: + self.set_step(updated_step) + + def __get__(self, obj, objtype=None): + if obj is None: + return self + self.update() + return self._x + + def __set__(self, obj, value): + self._x = value + + def __call__(self, step): + raise NotImplementedError diff --git a/deeplay/schedulers/sequence.py b/deeplay/schedulers/sequence.py new file mode 100644 index 00000000..4a3b25ce --- /dev/null +++ b/deeplay/schedulers/sequence.py @@ -0,0 +1,36 @@ +from .scheduler import BaseScheduler + + +class SchedulerSequence(BaseScheduler): + """Scheduler that returns value from one of the schedulers in the chain. + + The scheduler is chosen based on the current step. + """ + + def __init__(self, on_epoch=False): + super().__init__(on_epoch) + self.schedulers = [] + + def add(self, scheduler, n_steps=None): + if n_steps is None: + assert hasattr( + scheduler, "n_steps" + ), "For a scheduler without n_steps, n_steps must be specified" + n_steps = scheduler.n_steps + + self.schedulers.append((n_steps, scheduler)) + + def __call__(self, step): + for n_steps, scheduler in self.schedulers: + if step < n_steps: + return scheduler(step) + step -= n_steps + + final_step, final_scheduler = self.schedulers[-1] + return final_scheduler(final_step + step) + + def __repr__(self): + return f"{self.__class__.__name__}({self.schedulers})" + + def __str__(self): + return repr(self) diff --git a/deeplay/tests/schedulers/__init__.py b/deeplay/tests/schedulers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/deeplay/tests/schedulers/test_constant.py b/deeplay/tests/schedulers/test_constant.py new file mode 100644 index 00000000..d9d3652c --- /dev/null +++ b/deeplay/tests/schedulers/test_constant.py @@ -0,0 +1,89 @@ +import unittest + +import torch +from deeplay.applications.application import Application +from deeplay.components.mlp import MultiLayerPerceptron +from deeplay.external.optimizers.adam import Adam +from deeplay.module import DeeplayModule +from deeplay.schedulers import ConstantScheduler + + +class TestConstantScheduler(unittest.TestCase): + + def test_scheduler_build(self): + scheduler = ConstantScheduler(1.0) + scheduler.build() + + self.assertEqual(scheduler._step, 0) + self.assertIsNone(scheduler._x) + + def test_scheduler_step(self): + scheduler = ConstantScheduler(1.0) + scheduler.build() + + steps = [-999999999, -1, 0, 999, 9999999999999] + for step in steps: + value = scheduler(step) + self.assertEqual(value, 1.0) + + def test_scheduler_attached_to_module(self): + + class Module(DeeplayModule): + def __init__(self): + super().__init__() + self.x = ConstantScheduler(1.0) + + module = Module() + + # Before build, x is the scheduler + self.assertIsInstance(module.x, ConstantScheduler) + + module.build() + + # After build, x is the value of the scheduler + self.assertEqual(module.x, 1.0) + + def test_scheduler_attached_configure(self): + + class Module(DeeplayModule): + def __init__(self): + super().__init__() + self.x = ConstantScheduler(1.0) + + module = Module() + + # Before build, x is the scheduler + self.assertIsInstance(module.x, ConstantScheduler) + + module.x.configure(value=2.0) + + module.build() + + # After build, x is the value of the scheduler + self.assertEqual(module.x, 2.0) + + def test_scheduler_trainer(self): + + class Module(Application): + def __init__(self): + super().__init__(optimizer=Adam(lr=1.0), loss=torch.nn.MSELoss()) + self.x = ConstantScheduler(1.0) + self.net = MultiLayerPerceptron(1, [1], 1) + + def forward(_self, x): + self.assertEqual(_self.x, 1.0) + return _self.net(x) * _self.x + + module = Module() + + module.build() + + x = torch.randn(10, 1) + y = torch.randn(10, 1) + + module.fit((x, y), max_steps=10) + module._has_built = False + + self.assertEqual(module.x.trainer, module.trainer) + self.assertEqual(module.x._step, 9) + self.assertEqual(module.x._x, 1.0) diff --git a/deeplay/tests/schedulers/test_linear.py b/deeplay/tests/schedulers/test_linear.py new file mode 100644 index 00000000..b846185a --- /dev/null +++ b/deeplay/tests/schedulers/test_linear.py @@ -0,0 +1,120 @@ +import unittest + +import torch +from deeplay.applications.application import Application +from deeplay.components.mlp import MultiLayerPerceptron +from deeplay.external.optimizers.adam import Adam +from deeplay.module import DeeplayModule +from deeplay.schedulers import LinearScheduler + + +class TestConstantScheduler(unittest.TestCase): + + def test_scheduler_build(self): + scheduler = LinearScheduler(0.0, 1.0, 10) + scheduler.build() + + self.assertEqual(scheduler._step, 0) + self.assertIsNone(scheduler._x) + + def test_scheduler_step(self): + scheduler = LinearScheduler(0, 1, 10) + scheduler.build() + + steps = [-1, 0, 1, 9, 10, 20] + exp_values = [0.0, 0.0, 0.1, 0.9, 1.0, 1.0] + for step, exp_value in zip(steps, exp_values): + value = scheduler(step) + self.assertEqual(value, exp_value) + + def test_scheduler_step_negative(self): + scheduler = LinearScheduler(0, -1, 10) + scheduler.build() + + steps = [-1, 0, 1, 9, 10, 20] + exp_values = [0.0, 0.0, -0.1, -0.9, -1.0, -1.0] + for step, exp_value in zip(steps, exp_values): + value = scheduler(step) + self.assertEqual(value, exp_value) + + def test_scheduler_step_constant(self): + scheduler = LinearScheduler(0, 0, 10) + scheduler.build() + + steps = [-1, 0, 1, 9, 10, 20] + exp_values = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + for step, exp_value in zip(steps, exp_values): + value = scheduler(step) + self.assertEqual(value, exp_value) + + def test_scheduler_step_zero_steps(self): + scheduler = LinearScheduler(0, 1, 0) + scheduler.build() + + steps = [-1, 0, 1, 9, 10, 20] + exp_values = [0.0, 1.0, 1.0, 1.0, 1.0, 1.0] + for step, exp_value in zip(steps, exp_values): + value = scheduler(step) + self.assertEqual(value, exp_value) + + def test_scheduler_attached_to_module(self): + + class Module(DeeplayModule): + def __init__(self): + super().__init__() + self.x = LinearScheduler(0.0, 1.0, 10) + + module = Module() + + # Before build, x is the scheduler + self.assertIsInstance(module.x, LinearScheduler) + + module.build() + + # After build, x is the value of the scheduler + self.assertEqual(module.x, 0.0) + + def test_scheduler_attached_configure(self): + + class Module(DeeplayModule): + def __init__(self): + super().__init__() + self.x = LinearScheduler(0.0, 1.0, 10) + + module = Module() + + # Before build, x is the scheduler + self.assertIsInstance(module.x, LinearScheduler) + + module.x.configure(start_value=1.0, end_value=2.0) + + module.build() + + # After build, x is the value of the scheduler + self.assertEqual(module.x, 1.0) + + def test_scheduler_trainer(self): + + class Module(Application): + def __init__(self): + super().__init__(optimizer=Adam(lr=1.0), loss=torch.nn.MSELoss()) + self.x = LinearScheduler(0.0, 1.0, 10) + self.net = MultiLayerPerceptron(1, [1], 1) + + def forward(_self, x): + self.assertEqual(_self.x, _self.trainer.global_step / 10) + return _self.net(x) * _self.x + + module = Module() + + module.build() + + x = torch.randn(10, 1) + y = torch.randn(10, 1) + + module.fit((x, y), max_steps=10) + module._has_built = False + + self.assertEqual(module.x.trainer, module.trainer) + self.assertEqual(module.x._step, 9) + self.assertEqual(module.x._x, 0.9) diff --git a/deeplay/tests/schedulers/test_scheduler.py b/deeplay/tests/schedulers/test_scheduler.py new file mode 100644 index 00000000..7439ab1b --- /dev/null +++ b/deeplay/tests/schedulers/test_scheduler.py @@ -0,0 +1,12 @@ +import unittest +from deeplay.schedulers import BaseScheduler + + +class TestBaseScheduler(unittest.TestCase): + + def test_scheduler_defaults(self): + scheduler = BaseScheduler() + scheduler.build() + + self.assertEqual(scheduler._step, 0) + self.assertIsNone(scheduler._x) diff --git a/deeplay/tests/test_decorators.py b/deeplay/tests/test_decorators.py index ccc0c5b6..7f5de13a 100644 --- a/deeplay/tests/test_decorators.py +++ b/deeplay/tests/test_decorators.py @@ -63,8 +63,6 @@ def __init__(self): # module["encoder"] -# print("after:", module.encoder.p) - class DummyClass: def __init__(self):