diff --git a/CHANGELOG.md b/CHANGELOG.md index 3419199ac6..05feecdff1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Fixed a bug where the backbone learning rate would be divided by 10 when unfrozen if using the `freeze_unfreeze` or `unfreeze_milestones` strategies ([#1329](https://github.com/PyTorchLightning/lightning-flash/pull/1329)) +- Fixed naming of optimizer and scheduler registries which did not allow manual optimization. ([#1342](https://github.com/PyTorchLightning/lightning-flash/pull/1342)) + ## [0.7.4] - 2022-04-27 ### Fixed diff --git a/README.md b/README.md index 136478cbb8..235bb4b642 100644 --- a/README.md +++ b/README.md @@ -209,7 +209,7 @@ model = ImageClassifier(backbone="resnet18", num_classes=2, optimizer="Adam", lr You can also register you own custom scheduler recipes beforeahand and use them shown as above: ```py -@ImageClassifier.lr_schedulers +@ImageClassifier.lr_schedulers_registry def my_steplr_recipe(optimizer): return torch.optim.lr_scheduler.StepLR(optimizer, step_size=10) diff --git a/docs/source/general/optimization.rst b/docs/source/general/optimization.rst index 69e1ff13c0..ce17a2531f 100644 --- a/docs/source/general/optimization.rst +++ b/docs/source/general/optimization.rst @@ -176,7 +176,7 @@ Using the ``optimizers`` and ``lr_schedulers`` decorator pertaining to each :cla >>> import torch >>> from flash.image import ImageClassifier - >>> @ImageClassifier.lr_schedulers + >>> @ImageClassifier.lr_schedulers_registry ... def my_flash_steplr_recipe(optimizer): ... return torch.optim.lr_scheduler.StepLR(optimizer, step_size=10) ... diff --git a/flash/core/model.py b/flash/core/model.py index 47722ec3a5..0237500ab9 100644 --- a/flash/core/model.py +++ b/flash/core/model.py @@ -324,8 +324,8 @@ class Task(DatasetProcessor, ModuleWrapperBase, LightningModule, FineTuningHooks task. """ - optimizers: FlashRegistry = _OPTIMIZERS_REGISTRY - lr_schedulers: FlashRegistry = _SCHEDULERS_REGISTRY + optimizers_registry: FlashRegistry = _OPTIMIZERS_REGISTRY + lr_schedulers_registry: FlashRegistry = _SCHEDULERS_REGISTRY finetuning_strategies: FlashRegistry = _FINETUNING_STRATEGIES_REGISTRY outputs: FlashRegistry = BASE_OUTPUTS @@ -490,7 +490,7 @@ def _get_optimizer_class_from_registry(self, optimizer_key: str) -> Optimizer: f"\nUse `{self.__class__.__name__}.available_optimizers()` to list the available optimizers." f"\nList of available Optimizers: {self.available_optimizers()}." ) - optimizer_fn = self.optimizers.get(optimizer_key.lower()) + optimizer_fn = self.optimizers_registry.get(optimizer_key.lower()) return optimizer_fn def configure_optimizers(self) -> Union[Optimizer, Tuple[List[Optimizer], List[_LRScheduler]]]: @@ -614,7 +614,7 @@ def get_backbone_details(cls, key) -> List[str]: @classmethod def available_optimizers(cls) -> List[str]: """Returns a list containing the keys of the available Optimizers.""" - registry: Optional[FlashRegistry] = getattr(cls, "optimizers", None) + registry: Optional[FlashRegistry] = getattr(cls, "optimizers_registry", None) if registry is None: return [] return registry.available_keys() @@ -622,7 +622,7 @@ def available_optimizers(cls) -> List[str]: @classmethod def available_lr_schedulers(cls) -> List[str]: """Returns a list containing the keys of the available LR schedulers.""" - registry: Optional[FlashRegistry] = getattr(cls, "lr_schedulers", None) + registry: Optional[FlashRegistry] = getattr(cls, "lr_schedulers_registry", None) if registry is None: return [] return registry.available_keys() @@ -683,7 +683,7 @@ def _get_lr_scheduler_class_from_registry(self, lr_scheduler_key: str) -> Dict[s f"\nUse `{self.__class__.__name__}.available_lr_schedulers()` to list the available schedulers." f"\n>>> List of available LR Schedulers: {self.available_lr_schedulers()}." ) - lr_scheduler_fn: Dict[str, Any] = self.lr_schedulers.get(lr_scheduler_key.lower(), with_metadata=True) + lr_scheduler_fn: Dict[str, Any] = self.lr_schedulers_registry.get(lr_scheduler_key.lower(), with_metadata=True) return deepcopy(lr_scheduler_fn) def _instantiate_lr_scheduler(self, optimizer: Optimizer) -> Dict[str, Any]: diff --git a/tests/core/test_model.py b/tests/core/test_model.py index 0d3342976f..6c586e8ae9 100644 --- a/tests/core/test_model.py +++ b/tests/core/test_model.py @@ -307,12 +307,12 @@ def test_available_backbones_raises(): _ = ImageClassifier.available_backbones() -@ClassificationTask.lr_schedulers +@ClassificationTask.lr_schedulers_registry def custom_steplr_configuration_return_as_instance(optimizer): return torch.optim.lr_scheduler.StepLR(optimizer, step_size=10) -@ClassificationTask.lr_schedulers +@ClassificationTask.lr_schedulers_registry def custom_steplr_configuration_return_as_dict(optimizer): return { "scheduler": torch.optim.lr_scheduler.StepLR(optimizer, step_size=10), @@ -370,7 +370,7 @@ def test_optimizers_and_schedulers(tmpdir, optim, sched, interval): @pytest.mark.skipif(not _CORE_TESTING, reason="Not testing core.") def test_optimizer_learning_rate(): mock_optimizer = MagicMock() - Task.optimizers(mock_optimizer, "test") + Task.optimizers_registry(mock_optimizer, "test") model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax()) @@ -444,6 +444,36 @@ def train_dataloader(self): assert isinstance(trainer.lr_schedulers[0]["scheduler"], torch.optim.lr_scheduler.LambdaLR) +@pytest.mark.skipif(not _CORE_TESTING, reason="Not testing core.") +def test_manual_optimization(tmpdir): + class ManualOptimizationTask(Task): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.automatic_optimization = False + + def training_step(self, batch: Any, batch_idx: int) -> Any: + optimizers = self.optimizers() + assert isinstance(optimizers, torch.optim.Optimizer) + optimizers.zero_grad() + + output = self.step(batch, batch_idx, self.train_metrics) + self.manual_backward(output["loss"]) + + optimizers.step() + + lr_schedulers = self.lr_schedulers() + assert isinstance(lr_schedulers, torch.optim.lr_scheduler._LRScheduler) + lr_schedulers.step() + + model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.Softmax()) + train_dl = DataLoader(DummyDataset()) + val_dl = DataLoader(DummyDataset()) + task = ManualOptimizationTask(model, loss_fn=F.nll_loss, lr_scheduler=("steplr", {"step_size": 1})) + + trainer = pl.Trainer(fast_dev_run=True, default_root_dir=tmpdir) + trainer.fit(task, train_dl, val_dl) + + @pytest.mark.skipif(not _CORE_TESTING, reason="Not testing core.") def test_errors_and_exceptions_optimizers_and_schedulers(): model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax())