Merge pull request #101 from KamitaniLab/fix-scheduler-behavior

Fix problematic behavior of optimizer/scheduler in FeatureInversionTask

Author: ShuntaroAoki

bdpy/recon/torch/modules/__init__.py

@@ -2,3 +2,4 @@
 from .generator import build_generator, BaseGenerator
 from .latent import ArbitraryLatent, BaseLatent
 from .critic import TargetNormalizedMSE, BaseCritic
+from .optimizer import build_optimizer_factory, build_scheduler_factory

bdpy/recon/torch/modules/optimizer.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from functools import partial
+from itertools import chain
+
+if TYPE_CHECKING:
+    from typing import Dict, Any, Tuple, Union, Iterable, Callable
+    from typing_extensions import TypeAlias
+    from torch import Tensor
+    import torch.optim as optim
+    from ..modules import BaseGenerator, BaseLatent
+
+    # NOTE: The definition of `_ParamsT` is the same as in `torch.optim.optimizer`
+    #       in torch>=2.2.0. We define it here for compatibility with older versions.
+    _ParamsT: TypeAlias = Union[
+        Iterable[Tensor], Iterable[Dict[str, Any]], Iterable[Tuple[str, Tensor]]
+    ]
+
+    _OptimizerFactoryType: TypeAlias = Callable[
+        [BaseGenerator, BaseLatent], optim.Optimizer
+    ]
+    _SchedulerFactoryType: TypeAlias = Callable[
+        [optim.Optimizer], optim.lr_scheduler.LRScheduler
+    ]
+    _GetParamsFnType: TypeAlias = Callable[[BaseGenerator, BaseLatent], _ParamsT]
+
+
+def build_optimizer_factory(
+    optimizer_class: type[optim.Optimizer],
+    *,
+    get_params_fn: _GetParamsFnType | None = None,
+    **kwargs,
+) -> _OptimizerFactoryType:
+    """Build an optimizer factory.
+
+    Parameters
+    ----------
+    optimizer_class : type
+        Optimizer class.
+    get_params_fn : Callable[[BaseGenerator, BaseLatent], _ParamsT] | None
+        Custom function to get parameters from the generator and the latent.
+        If None, it uses `chain(generator.parameters(), latent.parameters())`.
+    kwargs : dict
+        Keyword arguments for the optimizer.
+
+    Returns
+    -------
+    Callable[[BaseGenerator, BaseLatent], optim.Optimizer]
+        Optimizer factory.
+
+    Examples
+    --------
+    >>> from torch.optim import Adam
+    >>> from bdpy.recon.torch.modules import build_optimizer_factory
+    >>> optimizer_factory = build_optimizer_factory(Adam, lr=1e-3)
+    >>> optimizer = optimizer_factory(generator, latent)
+    """
+    if get_params_fn is None:
+        get_params_fn = lambda generator, latent: chain(
+            generator.parameters(), latent.parameters()
+        )
+
+    def init_fn(generator: BaseGenerator, latent: BaseLatent) -> optim.Optimizer:
+        return optimizer_class(get_params_fn(generator, latent), **kwargs)
+
+    return init_fn
+
+
+def build_scheduler_factory(
+    scheduler_class: type[optim.lr_scheduler.LRScheduler], **kwargs
+) -> _SchedulerFactoryType:
+    """Build a scheduler factory.
+
+    Parameters
+    ----------
+    scheduler_class : type
+        Scheduler class.
+    kwargs : dict
+        Keyword arguments for the scheduler.
+
+    Returns
+    -------
+    Callable[[optim.Optimizer], optim.lr_scheduler.LRScheduler]
+        Scheduler factory.
+
+    Examples
+    --------
+    >>> from torch.optim.lr_scheduler import StepLR
+    >>> from bdpy.recon.torch.modules import build_scheduler_factory
+    >>> scheduler_factory = build_scheduler_factory(StepLR, step_size=100, gamma=0.1)
+    >>> scheduler = scheduler_factory(optimizer)
+    """
+    return partial(scheduler_class, **kwargs)

bdpy/recon/torch/task/inversion.py

@@ -1,21 +1,24 @@
 from __future__ import annotations
 
-from typing import Dict, Iterable, Callable
+from typing import Dict, Iterable, Callable, TYPE_CHECKING
 
 from itertools import chain
 
-import torch
-
-from ..modules import BaseEncoder, BaseGenerator, BaseLatent, BaseCritic
 from bdpy.task import BaseTask
 from bdpy.task.callback import BaseCallback, unused, _validate_callback
 
-FeatureType = Dict[str, torch.Tensor]
+if TYPE_CHECKING:
+    import torch
+
+    from ..modules import BaseEncoder, BaseGenerator, BaseLatent, BaseCritic
+    from ..modules.optimizer import _OptimizerFactoryType, _SchedulerFactoryType
+
+    _FeatureType = Dict[str, torch.Tensor]
 
 
 def _apply_to_features(
-    fn: Callable[[torch.Tensor], torch.Tensor], features: FeatureType
-) -> FeatureType:
+    fn: Callable[[torch.Tensor], torch.Tensor], features: _FeatureType
+) -> _FeatureType:
     return {k: fn(v) for k, v in features.items()}
 
 
@@ -115,10 +118,10 @@ class FeatureInversionTask(BaseTask):
         Latent variable module.
     critic : BaseCritic
         Critic module.
-    optimizer : torch.optim.Optimizer
-        Optimizer.
-    scheduler : torch.optim.lr_scheduler.LRScheduler, optional
-        Learning rate scheduler, by default None.
+    optimizer_factory : _OptimizerFactoryType
+        Factory function for optimizer.
+    scheduler_factory : _SchedulerFactoryType | None, optional
+        Factory function for scheduler, by default None.
     num_iterations : int, optional
         Number of iterations, by default 1.
     callbacks : FeatureInversionCallback | Iterable[FeatureInversionCallback] | None, optional
@@ -135,9 +138,9 @@ class FeatureInversionTask(BaseTask):
     >>> generator = build_generator(...)
     >>> latent = ArbitraryLatent(...)
     >>> critic = TargetNormalizedMSE(...)
-    >>> optimizer = torch.optim.Adam(latent.parameters())
+    >>> optimizer_factory = build_optimizer_factory(...)
     >>> task = FeatureInversionTask(
-    ...     encoder, generator, latent, critic, optimizer, num_iterations=200,
+    ...     encoder, generator, latent, critic, optimizer_factory, num_iterations=200,
     ... )
     >>> target_features = encoder(target_image)
     >>> reconstructed_image = task(target_features)
@@ -149,8 +152,8 @@ def __init__(
         generator: BaseGenerator,
         latent: BaseLatent,
         critic: BaseCritic,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler = None,
+        optimizer_factory: _OptimizerFactoryType,
+        scheduler_factory: _SchedulerFactoryType | None = None,
         num_iterations: int = 1,
         callbacks: FeatureInversionCallback
         | Iterable[FeatureInversionCallback]
@@ -161,14 +164,14 @@ def __init__(
         self._generator = generator
         self._latent = latent
         self._critic = critic
-        self._optimizer = optimizer
-        self._scheduler = scheduler
+        self._optimizer_factory = optimizer_factory
+        self._scheduler_factory = scheduler_factory
 
         self._num_iterations = num_iterations
 
     def run(
         self,
-        target_features: FeatureType,
+        target_features: _FeatureType,
     ) -> torch.Tensor:
         """Run feature inversion given target features.
 
@@ -217,10 +220,8 @@ def reset_states(self) -> None:
         """Reset the state of the task."""
         self._generator.reset_states()
         self._latent.reset_states()
-        self._optimizer = self._optimizer.__class__(
-            chain(
-                self._generator.parameters(),
-                self._latent.parameters(),
-            ),
-            **self._optimizer.defaults,
-        )
+        self._optimizer = self._optimizer_factory(self._generator, self._latent)
+        if self._scheduler_factory is not None:
+            self._scheduler = self._scheduler_factory(self._optimizer)
+        else:
+            self._scheduler = None

tests/recon/torch/modules/test_optimizer.py

@@ -0,0 +1,113 @@
+"""Tests for bdpy.recon.torch.modules.optimizer"""
+
+from __future__ import annotations
+
+import unittest
+
+from functools import partial
+import numpy as np
+import torch.nn as nn
+import torch.optim as optim
+from bdpy.recon.torch.modules import build_generator, ArbitraryLatent
+from bdpy.recon.torch.modules import build_optimizer_factory, build_scheduler_factory
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.fc = nn.Linear(in_dim, out_dim)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+class TestBuildOptimizerFactory(unittest.TestCase):
+    """Tests for bdpy.recon.torch.modules.optimizer.build_optimizer_factory"""
+
+    def test_build_optimizer_factory(self):
+        generator = build_generator(MLP(64, 10))
+        latent = ArbitraryLatent(
+            (1, 64), init_fn=partial(nn.init.normal_, mean=0, std=1)
+        )
+        optimizer_factory = build_optimizer_factory(optim.SGD, lr=0.1)
+        optimizer = optimizer_factory(generator, latent)
+        self.assertIsInstance(
+            optimizer,
+            optim.SGD,
+            msg="optimizer_factory should return an instance of optim.Optimizer",
+        )
+
+        latent.reset_states()
+        generator.reset_states()
+        latent_prev = latent().detach().clone().numpy()
+        optimizer.zero_grad()
+        output = generator(latent())
+        loss = output.sum()
+        loss.backward()
+        latent_next_expected = (
+            latent_prev - 0.1 * latent().grad.detach().clone().numpy()
+        )
+        optimizer.step()
+        latent_next = latent().detach().clone().numpy()
+        np.testing.assert_allclose(
+            latent_next,
+            latent_next_expected,
+            rtol=1e-6,
+            err_msg="Optimizer does not update the latent variable correctly.",
+        )
+
+        # check if all the frozen generator's gradients are None
+        generator_grad = [p.grad for p in generator.parameters()]
+        self.assertTrue(
+            all([g is None for g in generator_grad]),
+            msg="Frozen generator's gradients should be None after the optimizer step.",
+        )
+
+
+class TestBuildSchedulerFactory(unittest.TestCase):
+    """Tests for bdpy.recon.torch.modules.optimizer.build_scheduler_factory"""
+
+    def test_build_scheduler_factory(self):
+        generator = build_generator(MLP(64, 10))
+        latent = ArbitraryLatent(
+            (1, 64), init_fn=partial(nn.init.normal_, mean=0, std=1)
+        )
+        optimizer_factory = build_optimizer_factory(optim.SGD, lr=0.1)
+        scheduler_factory = build_scheduler_factory(
+            optim.lr_scheduler.StepLR, step_size=1, gamma=0.1
+        )
+        optimizer = optimizer_factory(generator, latent)
+        scheduler = scheduler_factory(optimizer)
+        self.assertIsInstance(
+            scheduler,
+            optim.lr_scheduler.StepLR,
+            msg="Scheduler factory should return an instance of optim.lr_scheduler.LRScheduler",
+        )
+
+        latent.reset_states()
+        generator.reset_states()
+        optimizer.zero_grad()
+        output = generator(latent())
+        loss = output.sum()
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+        self.assertEqual(
+            optimizer.param_groups[0]["lr"],
+            0.1 * 0.1,
+            "Scheduler does not update the learning rate correctly.",
+        )
+
+        # check if reference to the optimizer is kept during re-initialization
+        for _ in range(10):
+            optimizer = optimizer_factory(generator, latent)
+            scheduler = scheduler_factory(optimizer)
+        else:
+            self.assertTrue(
+                scheduler.optimizer is optimizer,
+                "Scheduler should keep the reference to the optimizer during re-initialization.",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()